diff --git a/.gitattributes b/.gitattributes index a6344aac8c09253b3b630fb776ae94478aa0275b..0cd58331b2a989b68be4ec5676383437fca8687b 100644 --- a/.gitattributes +++ b/.gitattributes @@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text *.zip filter=lfs diff=lfs merge=lfs -text *.zst filter=lfs diff=lfs merge=lfs -text *tfevents* filter=lfs diff=lfs merge=lfs -text +*.so filter=lfs diff=lfs merge=lfs -text diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..afb3b993a1e53b36066133d5ee9de3f3a513fed7 --- /dev/null +++ b/.gitignore @@ -0,0 +1,3 @@ +.bak +__pycache__ +result \ No newline at end of file diff --git a/.make_markers/patch_applied b/.make_markers/patch_applied new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/.make_markers/submodule_initialized b/.make_markers/submodule_initialized new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/README.md b/README.md new file mode 100644 index 0000000000000000000000000000000000000000..8029d9b21c46c6ae82925ddadd8e007b491dd7eb --- /dev/null +++ b/README.md @@ -0,0 +1,10 @@ +--- +license: apache-2.0 +tags: + - kernel +--- + + +This kernel is a work in progress and requires more work to correctly add all of the FlashInfer kernels. + +Please see the [generate-source.md](generate-source.md) for instructions on how to generate the source files that are contained in this kernel. \ No newline at end of file diff --git a/build.toml b/build.toml new file mode 100644 index 0000000000000000000000000000000000000000..d3ee110725ff3e487c9888c0f808628afff2bd6c --- /dev/null +++ b/build.toml @@ -0,0 +1,224 @@ +[general] +name = "flashinfer" +universal = false + +[torch] +src = [ + "torch-ext/torch_binding.cpp", + "torch-ext/torch_binding.h" +] + +[kernel.flashinfer] +backend = "cuda" +cuda-capabilities = [ + # "7.0", + # "7.2", + "7.5", + "8.0", + "8.6", + "8.7", + "8.9", + "9.0", + # "10.0", + # "10.1", + # "11.8", + # "12.0" +] +cuda-flags = [ + "-O3", + "-std=c++17", + "--use_fast_math", + "--expt-relaxed-constexpr", + "--expt-extended-lambda", + "-DFLASHINFER_ENABLE_F16", + "-DFLASHINFER_ENABLE_BF16", + "-DFLASHINFER_ENABLE_FP8_E4M3", + "-DFLASHINFER_ENABLE_FP8_E5M2", + "-DNDEBUG", +] +cxx-flags = [ + "-DFLASHINFER_ENABLE_F16", + "-DFLASHINFER_ENABLE_BF16", + "-DFLASHINFER_ENABLE_FP8_E4M3", + "-DFLASHINFER_ENABLE_FP8_E5M2", +] +include = [ "include" ] +depends = ["torch", "cutlass_3_8"] +src = [ + "include/flashinfer/page.cuh", + "include/flashinfer/arch_condition.h", + "include/flashinfer/cubin_loader.h", + "include/flashinfer/permuted_smem.cuh", + "include/flashinfer/exception.h", + "include/flashinfer/trtllm/common.h", + "include/flashinfer/trtllm/fmha/fmhaRunnerParams.h", + "include/flashinfer/trtllm/fmha/decoder_impl_common.h", + "include/flashinfer/trtllm/fmha/decoder_params.h", + "include/flashinfer/trtllm/fmha/kernelParams.h", + "include/flashinfer/trtllm/fmha/gen_kernel_launcher.cuh", + "include/flashinfer/trtllm/fmha/fmhaKernels.cuh", + "include/flashinfer/trtllm/fmha/cubin/kernelMetaInfo.h", + "include/flashinfer/trtllm/fmha/fmhaRunner.cuh", + "include/flashinfer/trtllm/common/cudaTypeUtils.cuh", + "include/flashinfer/trtllm/common/cudaUtils.h", + "include/flashinfer/trtllm/common/cudaBf16Wrapper.h", + "include/flashinfer/trtllm/common/cudaFp8Utils.h", + "include/flashinfer/trtllm/common/cudaBf16Fallbacks.cuh", + "include/flashinfer/trtllm/fused_moe/RoutingKernel.cuh", + "include/flashinfer/trtllm/fused_moe/DevKernel.h", + "include/flashinfer/trtllm/fused_moe/RoutingKernel.h", + "include/flashinfer/trtllm/fused_moe/runner.h", + "include/flashinfer/trtllm/fused_moe/IntFastDiv.h", + "include/flashinfer/trtllm/fused_moe/RoutingKernelTopK.cuh", + "include/flashinfer/trtllm/batched_gemm/KernelRunner.h", + "include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/GemmGatedActOptions.h", + "include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/MmaDecl.h", + "include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/SfLayoutDecl.h", + "include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/DtypeDecl.h", + "include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/CommonUtils.h", + "include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/CudaKernelLauncher.h", + "include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/BatchedGemmOptions.h", + "include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/TmaDescriptor.h", + "include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/BatchedGemmEnums.h", + "include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/KernelParamsDecl.h", + "include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/KernelTraits.h", + "include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/KernelParams.h", + "include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/Enums.h", + "include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/KernelMetaInfo.h", + "include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/BatchedGemmInterface.h", + "include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/GemmOptions.h", + "include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/trtllm/gen/MmaDecl.h", + "include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/trtllm/gen/SfLayoutDecl.h", + "include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/trtllm/gen/DtypeDecl.h", + "include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/trtllm/gen/CommonUtils.h", + "include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/trtllm/gen/CudaKernelLauncher.h", + "include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/TmaDescriptor.h", + "include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/KernelTraits.h", + "include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/GemmInterface.h", + "include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/KernelParams.h", + "include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/Enums.h", + "include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/GemmOptions.h", + "include/flashinfer/page.cuh", + "include/flashinfer/vec_dtypes.cuh", + "include/flashinfer/sampling.cuh", + "include/flashinfer/logging.h", + "include/flashinfer/fp16.h", + "include/flashinfer/attention_impl.cuh", + "include/flashinfer/allocator.h", + "include/flashinfer/utils.cuh", + "include/flashinfer/fastdiv.cuh", + "include/flashinfer/norm.cuh", + "include/flashinfer/layout.cuh", + "include/flashinfer/math.cuh", + "include/flashinfer/gemm/group_gemm_sm90.cuh", + "include/flashinfer/gemm/group_gemv.cuh", + "include/flashinfer/gemm/fp8_gemm_cutlass_template.h", + "include/flashinfer/gemm/group_gemm_mxfp4_groupwise_sm100.cuh", + "include/flashinfer/gemm/fp4_gemm_cutlass.h", + "include/flashinfer/gemm/fp4_gemm_cutlass_template.h", + "include/flashinfer/gemm/cutlass_gemm_configs.h", + "include/flashinfer/gemm/fp4_gemm_template_sm100.h", + "include/flashinfer/gemm/group_gemm_fp8_groupwise_sm100.cuh", + "include/flashinfer/gemm/gemm_groupwise_sm100.cuh", + "include/flashinfer/gemm/group_gemm_lora.cuh", + "include/flashinfer/gemm/group_gemm.cuh", + "include/flashinfer/gemm/fp8_gemm_cutlass.h", + "include/flashinfer/gemm/fp8_gemm_template_sm100.h", + "include/flashinfer/gemm/bmm_fp8.cuh", + "include/flashinfer/activation.cuh", + "include/flashinfer/semaphore_utils.cuh", + "include/flashinfer/quantization.cuh", + "include/flashinfer/mma.cuh", + "include/flashinfer/attention/variant_helper.cuh", + "include/flashinfer/attention/hopper.cuh", + "include/flashinfer/attention/prefill.cuh", + "include/flashinfer/attention/persistent_template.cuh", + "include/flashinfer/attention/mla_hopper.cuh", + "include/flashinfer/attention/pod.cuh", + "include/flashinfer/attention/variants.cuh", + "include/flashinfer/attention/decode.cuh", + "include/flashinfer/attention/heap.h", + "include/flashinfer/attention/state.cuh", + "include/flashinfer/attention/mask.cuh", + "include/flashinfer/attention/default_prefill_params.cuh", + "include/flashinfer/attention/cutlass_mla.cuh", + "include/flashinfer/attention/scheduler.cuh", + "include/flashinfer/attention/blackwell/kernel/fmha_options.hpp", + "include/flashinfer/attention/blackwell/kernel/gather_tensor.hpp", + "include/flashinfer/attention/blackwell/kernel/fmha_tile_scheduler.hpp", + "include/flashinfer/attention/blackwell/kernel/sm100_fmha_fwd_kernel_tma_warpspecialized.hpp", + "include/flashinfer/attention/blackwell/kernel/sm100_fmha_mla_tma_warpspecialized.hpp", + "include/flashinfer/attention/blackwell/kernel/sm100_mla_tile_scheduler.hpp", + "include/flashinfer/attention/blackwell/kernel/sm100_fmha_gen_kernel_warpspecialized.hpp", + "include/flashinfer/attention/blackwell/kernel/sm100_fmha_mla_reduction.hpp", + "include/flashinfer/attention/blackwell/plan.cuh", + "include/flashinfer/attention/blackwell/common/pow_2.hpp", + "include/flashinfer/attention/blackwell/collective/sm100_fmha_load_tma_warpspecialized.hpp", + "include/flashinfer/attention/blackwell/collective/sm100_fmha_fwd_epilogue_tma_warpspecialized.hpp", + "include/flashinfer/attention/blackwell/collective/sm100_fmha_gen_mainloop_warpspecialized.hpp", + "include/flashinfer/attention/blackwell/collective/sm100_fmha_load_cpasync_warpspecialized.hpp", + "include/flashinfer/attention/blackwell/collective/sm100_fmha_gen_epilogue_warpspecialized.hpp", + "include/flashinfer/attention/blackwell/collective/sm100_fmha_fwd_mainloop_tma_warpspecialized.hpp", + "include/flashinfer/attention/blackwell/collective/fmha_fusion.hpp", + "include/flashinfer/attention/blackwell/collective/fmha_common.hpp", + "include/flashinfer/attention/blackwell/device/sm100_mla.hpp", + "include/flashinfer/attention/blackwell/device/fmha.hpp", + "include/flashinfer/attention/blackwell/fmha_cutlass_sm100.cuh", + "include/flashinfer/attention/mla.cuh", + "include/flashinfer/attention/decode_mla_cute_sm80.cuh", + "include/flashinfer/attention/cascade.cuh", + "include/flashinfer/attention/hopper/variant_helper.cuh", + "include/flashinfer/attention/hopper/default_params.cuh", + "include/flashinfer/attention/hopper/attention_updater.cuh", + "include/flashinfer/attention/hopper/epilogue.cuh", + "include/flashinfer/attention/hopper/variants.cuh", + "include/flashinfer/attention/hopper/mainloop.cuh", + "include/flashinfer/attention/hopper/block_sparse_gather.cuh", + "include/flashinfer/attention/hopper/utils.cuh", + "include/flashinfer/attention/hopper/prefill_sm90.cuh", + "include/flashinfer/attention/hopper/named_barrier.cuh", + "include/flashinfer/attention/hopper/tile_scheduler.cuh", + "include/flashinfer/attention/hopper/sparse_mainloop.cuh", + "include/flashinfer/attention/hopper/kernel_traits.cuh", + "include/flashinfer/attention/hopper/mainloop_mma.cuh", + "include/flashinfer/attention/hopper/quantization/epilogue.cuh", + "include/flashinfer/attention/hopper/quantization/prefill_sm90.cuh", + "include/flashinfer/attention/hopper/quantization/mainloop_load.cuh", + "include/flashinfer/attention/hopper/quantization/mainloop_sparse_load.cuh", + "include/flashinfer/attention/hopper/quantization/kernel_traits.cuh", + "include/flashinfer/attention/hopper/quantization/mainloop_mma.cuh", + "include/flashinfer/attention/default_decode_params.cuh", + "include/flashinfer/attention/mla_params.cuh", + "include/flashinfer/attention/persistent.cuh", + "include/flashinfer/pos_enc.cuh", + "include/flashinfer/cutlass_utils.cuh", + "include/flashinfer/comm/trtllm_moe_allreduce_fusion.cuh", + "include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh", + "include/flashinfer/comm/vllm_custom_all_reduce.cuh", + "include/flashinfer/comm/trtllm_allreduce.cuh", + "include/flashinfer/comm/trtllm_allreduce_fusion.cuh", + "include/flashinfer/comm/trtllm_alltoall.cuh", + "include/flashinfer/frag_layout_swizzle.cuh", + "include/flashinfer/profiler.cuh", + "include/flashinfer/cp_async.cuh", + + "include/pytorch_conversion_utils.h", + "include/pytorch_extension_utils.h", + + # # batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False + # "csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_config.inc", + # "csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_jit_pybind.cu", + # "csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode.cu", + # "csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_kernel.cu", + + # # batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False + # "csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode_config.inc", + # "csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode_jit_pybind.cu", + # "csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode.cu", + # "csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode_kernel.cu", + + + "csrc/generated/gelu_and_mul.cu", + "csrc/generated/gelu_tanh_and_mul.cu", + "csrc/generated/silu_and_mul.cu", +] \ No newline at end of file diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode.cu b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode.cu new file mode 100644 index 0000000000000000000000000000000000000000..02c75e02718f8f8ce9b6036f8f788d9a9e254cf1 --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode.cu @@ -0,0 +1,197 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include +#include + +#include "batch_decode_config.inc" +#include "pytorch_conversion_utils.h" +#include "pytorch_extension_utils.h" + +namespace flashinfer { + +template +cudaError_t BatchDecodeWithPagedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v, + float* tmp_s, bool enable_pdl, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +at::Tensor BatchDecodeWithPagedKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size, + int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, + int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo, + at::Tensor empty_q_data, at::Tensor empty_kv_data) { + size_t float_workspace_size_in_bytes = + float_workspace_buffer.size(0) * float_workspace_buffer.element_size(); + size_t int_workspace_size_in_bytes = + int_workspace_buffer.size(0) * int_workspace_buffer.element_size(); + + DecodePlanInfo plan_info; + + auto q_scalar_type = empty_q_data.scalar_type(); + auto kv_scalar_type = empty_kv_data.scalar_type(); + + TORCH_CHECK(head_dim_qk == head_dim_vo, + "CUDA cores template only supports equal head dim for QK and VO, please use tensor " + "cores template for different head dim"); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + DISPATCH_GQA_GROUP_SIZE(num_qo_heads / num_kv_heads, GROUP_SIZE, { + auto work_estimation_func = BatchDecodeWithPagedKVCacheWorkEstimationDispatched< + GROUP_SIZE, HEAD_DIM_QK, POS_ENCODING_MODE, AttentionVariant, Params>; + cudaError_t status = DecodePlan( + static_cast(float_workspace_buffer.data_ptr()), float_workspace_size_in_bytes, + static_cast(int_workspace_buffer.data_ptr()), + static_cast(page_locked_int_workspace_buffer.data_ptr()), + int_workspace_size_in_bytes, plan_info, static_cast(indptr.data_ptr()), + batch_size, num_qo_heads, page_size, enable_cuda_graph, + /*stream=*/stream, work_estimation_func); + + TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); + }); + + return vec_to_tensor(plan_info.ToVector()); +} + +void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, + at::Tensor paged_v_cache, at::Tensor paged_kv_indptr, + at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, + int64_t kv_layout_code, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + DecodePlanInfo plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + QKVLayout kv_layout = static_cast(kv_layout_code); + auto device = q.device(); + int64_t batch_size = q.size(0); + int64_t num_qo_heads = q.size(1); + int64_t num_kv_heads, page_size; + + if (kv_layout == QKVLayout::kHND) { + num_kv_heads = paged_k_cache.size(1); + page_size = paged_k_cache.size(2); + } else { + page_size = paged_k_cache.size(1); + num_kv_heads = paged_k_cache.size(2); + } + uint32_t head_dim_qk = q.size(2); + uint32_t head_dim_vo = paged_v_cache.size(3); + + TORCH_CHECK(head_dim_qk == head_dim_vo, + "CUDA cores template only supports equal head dim for QK and VO, please use tensor " + "cores template for different head dim"); + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == batch_size, lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == num_qo_heads, lse.size(1), q.size(1)); + } + + void* float_buffer = static_cast(float_workspace_buffer.data_ptr()); + void* int_buffer = static_cast(int_workspace_buffer.data_ptr()); + + // get q_scalar_type and kv_scalar_type + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = paged_k_cache.scalar_type(); + + // get q_stride_n and q_stride_h + const auto q_stride_n = q.stride(0); + const auto q_stride_h = q.stride(1); + + // get kv_cache_strides + const int64_t* kv_cache_strides = nullptr; + auto k_strides = paged_k_cache.strides(); + auto v_strides = paged_v_cache.strides(); + TORCH_CHECK(k_strides == v_strides, "k/v strides must be identical"); + kv_cache_strides = k_strides.data(); + + const c10::cuda::OptionalCUDAGuard device_guard(device); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + paged_kv_t paged_kv( + num_kv_heads, page_size, HEAD_DIM_QK, batch_size, kv_layout, + static_cast(paged_k_cache.data_ptr()), + static_cast(paged_v_cache.data_ptr()), kv_cache_strides, + static_cast(paged_kv_indices.data_ptr()), + static_cast(paged_kv_indptr.data_ptr()), + static_cast(paged_kv_last_page_len.data_ptr())); + + Params params; + params.q = static_cast(q.data_ptr()); + params.paged_kv = paged_kv; + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.padded_batch_size = 0; + params.num_qo_heads = num_qo_heads; + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.window_left = window_left; + params.request_indices = nullptr; + params.kv_tile_indices = nullptr; + params.o_indptr = nullptr; + params.kv_chunk_size_ptr = nullptr; + params.block_valid_mask = nullptr; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + DTypeO* tmp_v = nullptr; + float* tmp_s = nullptr; + params.request_indices = + GetPtrFromBaseOffset(int_buffer, plan_info.request_indices_offset); + params.kv_tile_indices = + GetPtrFromBaseOffset(int_buffer, plan_info.kv_tile_indices_offset); + params.o_indptr = GetPtrFromBaseOffset(int_buffer, plan_info.o_indptr_offset); + params.kv_chunk_size_ptr = + GetPtrFromBaseOffset(int_buffer, plan_info.kv_chunk_size_ptr_offset); + if (plan_info.split_kv) { + tmp_v = GetPtrFromBaseOffset(float_buffer, plan_info.v_offset); + tmp_s = GetPtrFromBaseOffset(float_buffer, plan_info.s_offset); + if (plan_info.enable_cuda_graph) { + params.block_valid_mask = + GetPtrFromBaseOffset(int_buffer, plan_info.block_valid_mask_offset); + } + } + params.padded_batch_size = plan_info.padded_batch_size; + + cudaError_t status = + flashinfer::BatchDecodeWithPagedKVCacheDispatched(params, tmp_v, + tmp_s, enable_pdl, + /*stream=*/stream); + TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); +} diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_config.inc b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..1e47391fea886308ecb359c02d0e23df6ba31640 --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_config.inc @@ -0,0 +1,71 @@ +#pragma once +#include +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta +#define ADDITIONAL_PARAMS_SETTER params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, ...) { \ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ +} + +using namespace flashinfer; + +using DTypeQ = nv_bfloat16; +using DTypeKV = nv_bfloat16; +using DTypeO = nv_bfloat16; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 128; +constexpr int HEAD_DIM_VO = 128; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = false; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + DTypeQ* q; + paged_kv_t paged_kv; + DTypeO* o; + float* lse; + + float* maybe_alibi_slopes; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; + + + uint32_t padded_batch_size; + uint32_t num_qo_heads; + IdType q_stride_n; + IdType q_stride_h; + int32_t window_left; + bool enable_pdl; + + IdType* request_indices; + IdType* kv_tile_indices; + IdType* o_indptr; + IdType* kv_chunk_size_ptr; + bool* block_valid_mask; + bool partition_kv; + + __host__ __device__ __forceinline__ int32_t get_qo_len(int32_t batch_idx) const { return 1; } + + __host__ __device__ __forceinline__ int32_t get_kv_len(int32_t batch_idx) const { + return paged_kv.get_length(batch_idx); + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_jit_pybind.cu b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..91e944770b8cee57481f0631545cef02694f50cc --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_jit_pybind.cu @@ -0,0 +1,40 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "batch_decode_config.inc" +#include "pytorch_extension_utils.h" + +at::Tensor BatchDecodeWithPagedKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size, + int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, + int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo, + at::Tensor empty_q_data, at::Tensor empty_kv_data); + +void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, + at::Tensor paged_v_cache, at::Tensor paged_kv_indptr, + at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, + int64_t kv_layout_code, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Batched decode with paged KV-Cache plan + m.def("plan", BatchDecodeWithPagedKVCachePlan); + // Batched decode with paged KV-Cache run + m.def("run", BatchDecodeWithPagedKVCacheRun); +} diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_kernel.cu b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_kernel.cu new file mode 100644 index 0000000000000000000000000000000000000000..973142e154c74a17ca0e21d47d510e1167de4891 --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_kernel.cu @@ -0,0 +1,13 @@ +#include +#include "batch_decode_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t +BatchDecodeWithPagedKVCacheDispatched<128, PosEncodingMode::kNone, DefaultAttention, Params>( + Params params, nv_bfloat16* tmp_v, + float* tmp_s, bool enable_pdl, cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode.cu b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode.cu new file mode 100644 index 0000000000000000000000000000000000000000..02c75e02718f8f8ce9b6036f8f788d9a9e254cf1 --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode.cu @@ -0,0 +1,197 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include +#include + +#include "batch_decode_config.inc" +#include "pytorch_conversion_utils.h" +#include "pytorch_extension_utils.h" + +namespace flashinfer { + +template +cudaError_t BatchDecodeWithPagedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v, + float* tmp_s, bool enable_pdl, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +at::Tensor BatchDecodeWithPagedKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size, + int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, + int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo, + at::Tensor empty_q_data, at::Tensor empty_kv_data) { + size_t float_workspace_size_in_bytes = + float_workspace_buffer.size(0) * float_workspace_buffer.element_size(); + size_t int_workspace_size_in_bytes = + int_workspace_buffer.size(0) * int_workspace_buffer.element_size(); + + DecodePlanInfo plan_info; + + auto q_scalar_type = empty_q_data.scalar_type(); + auto kv_scalar_type = empty_kv_data.scalar_type(); + + TORCH_CHECK(head_dim_qk == head_dim_vo, + "CUDA cores template only supports equal head dim for QK and VO, please use tensor " + "cores template for different head dim"); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + DISPATCH_GQA_GROUP_SIZE(num_qo_heads / num_kv_heads, GROUP_SIZE, { + auto work_estimation_func = BatchDecodeWithPagedKVCacheWorkEstimationDispatched< + GROUP_SIZE, HEAD_DIM_QK, POS_ENCODING_MODE, AttentionVariant, Params>; + cudaError_t status = DecodePlan( + static_cast(float_workspace_buffer.data_ptr()), float_workspace_size_in_bytes, + static_cast(int_workspace_buffer.data_ptr()), + static_cast(page_locked_int_workspace_buffer.data_ptr()), + int_workspace_size_in_bytes, plan_info, static_cast(indptr.data_ptr()), + batch_size, num_qo_heads, page_size, enable_cuda_graph, + /*stream=*/stream, work_estimation_func); + + TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); + }); + + return vec_to_tensor(plan_info.ToVector()); +} + +void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, + at::Tensor paged_v_cache, at::Tensor paged_kv_indptr, + at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, + int64_t kv_layout_code, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + DecodePlanInfo plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + QKVLayout kv_layout = static_cast(kv_layout_code); + auto device = q.device(); + int64_t batch_size = q.size(0); + int64_t num_qo_heads = q.size(1); + int64_t num_kv_heads, page_size; + + if (kv_layout == QKVLayout::kHND) { + num_kv_heads = paged_k_cache.size(1); + page_size = paged_k_cache.size(2); + } else { + page_size = paged_k_cache.size(1); + num_kv_heads = paged_k_cache.size(2); + } + uint32_t head_dim_qk = q.size(2); + uint32_t head_dim_vo = paged_v_cache.size(3); + + TORCH_CHECK(head_dim_qk == head_dim_vo, + "CUDA cores template only supports equal head dim for QK and VO, please use tensor " + "cores template for different head dim"); + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == batch_size, lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == num_qo_heads, lse.size(1), q.size(1)); + } + + void* float_buffer = static_cast(float_workspace_buffer.data_ptr()); + void* int_buffer = static_cast(int_workspace_buffer.data_ptr()); + + // get q_scalar_type and kv_scalar_type + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = paged_k_cache.scalar_type(); + + // get q_stride_n and q_stride_h + const auto q_stride_n = q.stride(0); + const auto q_stride_h = q.stride(1); + + // get kv_cache_strides + const int64_t* kv_cache_strides = nullptr; + auto k_strides = paged_k_cache.strides(); + auto v_strides = paged_v_cache.strides(); + TORCH_CHECK(k_strides == v_strides, "k/v strides must be identical"); + kv_cache_strides = k_strides.data(); + + const c10::cuda::OptionalCUDAGuard device_guard(device); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + paged_kv_t paged_kv( + num_kv_heads, page_size, HEAD_DIM_QK, batch_size, kv_layout, + static_cast(paged_k_cache.data_ptr()), + static_cast(paged_v_cache.data_ptr()), kv_cache_strides, + static_cast(paged_kv_indices.data_ptr()), + static_cast(paged_kv_indptr.data_ptr()), + static_cast(paged_kv_last_page_len.data_ptr())); + + Params params; + params.q = static_cast(q.data_ptr()); + params.paged_kv = paged_kv; + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.padded_batch_size = 0; + params.num_qo_heads = num_qo_heads; + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.window_left = window_left; + params.request_indices = nullptr; + params.kv_tile_indices = nullptr; + params.o_indptr = nullptr; + params.kv_chunk_size_ptr = nullptr; + params.block_valid_mask = nullptr; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + DTypeO* tmp_v = nullptr; + float* tmp_s = nullptr; + params.request_indices = + GetPtrFromBaseOffset(int_buffer, plan_info.request_indices_offset); + params.kv_tile_indices = + GetPtrFromBaseOffset(int_buffer, plan_info.kv_tile_indices_offset); + params.o_indptr = GetPtrFromBaseOffset(int_buffer, plan_info.o_indptr_offset); + params.kv_chunk_size_ptr = + GetPtrFromBaseOffset(int_buffer, plan_info.kv_chunk_size_ptr_offset); + if (plan_info.split_kv) { + tmp_v = GetPtrFromBaseOffset(float_buffer, plan_info.v_offset); + tmp_s = GetPtrFromBaseOffset(float_buffer, plan_info.s_offset); + if (plan_info.enable_cuda_graph) { + params.block_valid_mask = + GetPtrFromBaseOffset(int_buffer, plan_info.block_valid_mask_offset); + } + } + params.padded_batch_size = plan_info.padded_batch_size; + + cudaError_t status = + flashinfer::BatchDecodeWithPagedKVCacheDispatched(params, tmp_v, + tmp_s, enable_pdl, + /*stream=*/stream); + TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); +} diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode_config.inc b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..6919f34bcddd30e2293290f4db7eb3fd71a73768 --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode_config.inc @@ -0,0 +1,71 @@ +#pragma once +#include +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta +#define ADDITIONAL_PARAMS_SETTER params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, ...) { \ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ +} + +using namespace flashinfer; + +using DTypeQ = nv_bfloat16; +using DTypeKV = nv_bfloat16; +using DTypeO = nv_bfloat16; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 256; +constexpr int HEAD_DIM_VO = 256; +constexpr auto USE_LOGITS_SOFT_CAP = true; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = true; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + DTypeQ* q; + paged_kv_t paged_kv; + DTypeO* o; + float* lse; + + float* maybe_alibi_slopes; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; + + + uint32_t padded_batch_size; + uint32_t num_qo_heads; + IdType q_stride_n; + IdType q_stride_h; + int32_t window_left; + bool enable_pdl; + + IdType* request_indices; + IdType* kv_tile_indices; + IdType* o_indptr; + IdType* kv_chunk_size_ptr; + bool* block_valid_mask; + bool partition_kv; + + __host__ __device__ __forceinline__ int32_t get_qo_len(int32_t batch_idx) const { return 1; } + + __host__ __device__ __forceinline__ int32_t get_kv_len(int32_t batch_idx) const { + return paged_kv.get_length(batch_idx); + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode_jit_pybind.cu b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..91e944770b8cee57481f0631545cef02694f50cc --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode_jit_pybind.cu @@ -0,0 +1,40 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "batch_decode_config.inc" +#include "pytorch_extension_utils.h" + +at::Tensor BatchDecodeWithPagedKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size, + int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, + int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo, + at::Tensor empty_q_data, at::Tensor empty_kv_data); + +void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, + at::Tensor paged_v_cache, at::Tensor paged_kv_indptr, + at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, + int64_t kv_layout_code, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Batched decode with paged KV-Cache plan + m.def("plan", BatchDecodeWithPagedKVCachePlan); + // Batched decode with paged KV-Cache run + m.def("run", BatchDecodeWithPagedKVCacheRun); +} diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode_kernel.cu b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode_kernel.cu new file mode 100644 index 0000000000000000000000000000000000000000..b63da6a66b8eeb551357ef21128a476cd6477474 --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode_kernel.cu @@ -0,0 +1,13 @@ +#include +#include "batch_decode_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t +BatchDecodeWithPagedKVCacheDispatched<256, PosEncodingMode::kNone, DefaultAttention, Params>( + Params params, nv_bfloat16* tmp_v, + float* tmp_s, bool enable_pdl, cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode.cu b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode.cu new file mode 100644 index 0000000000000000000000000000000000000000..02c75e02718f8f8ce9b6036f8f788d9a9e254cf1 --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode.cu @@ -0,0 +1,197 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include +#include + +#include "batch_decode_config.inc" +#include "pytorch_conversion_utils.h" +#include "pytorch_extension_utils.h" + +namespace flashinfer { + +template +cudaError_t BatchDecodeWithPagedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v, + float* tmp_s, bool enable_pdl, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +at::Tensor BatchDecodeWithPagedKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size, + int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, + int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo, + at::Tensor empty_q_data, at::Tensor empty_kv_data) { + size_t float_workspace_size_in_bytes = + float_workspace_buffer.size(0) * float_workspace_buffer.element_size(); + size_t int_workspace_size_in_bytes = + int_workspace_buffer.size(0) * int_workspace_buffer.element_size(); + + DecodePlanInfo plan_info; + + auto q_scalar_type = empty_q_data.scalar_type(); + auto kv_scalar_type = empty_kv_data.scalar_type(); + + TORCH_CHECK(head_dim_qk == head_dim_vo, + "CUDA cores template only supports equal head dim for QK and VO, please use tensor " + "cores template for different head dim"); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + DISPATCH_GQA_GROUP_SIZE(num_qo_heads / num_kv_heads, GROUP_SIZE, { + auto work_estimation_func = BatchDecodeWithPagedKVCacheWorkEstimationDispatched< + GROUP_SIZE, HEAD_DIM_QK, POS_ENCODING_MODE, AttentionVariant, Params>; + cudaError_t status = DecodePlan( + static_cast(float_workspace_buffer.data_ptr()), float_workspace_size_in_bytes, + static_cast(int_workspace_buffer.data_ptr()), + static_cast(page_locked_int_workspace_buffer.data_ptr()), + int_workspace_size_in_bytes, plan_info, static_cast(indptr.data_ptr()), + batch_size, num_qo_heads, page_size, enable_cuda_graph, + /*stream=*/stream, work_estimation_func); + + TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); + }); + + return vec_to_tensor(plan_info.ToVector()); +} + +void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, + at::Tensor paged_v_cache, at::Tensor paged_kv_indptr, + at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, + int64_t kv_layout_code, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + DecodePlanInfo plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + QKVLayout kv_layout = static_cast(kv_layout_code); + auto device = q.device(); + int64_t batch_size = q.size(0); + int64_t num_qo_heads = q.size(1); + int64_t num_kv_heads, page_size; + + if (kv_layout == QKVLayout::kHND) { + num_kv_heads = paged_k_cache.size(1); + page_size = paged_k_cache.size(2); + } else { + page_size = paged_k_cache.size(1); + num_kv_heads = paged_k_cache.size(2); + } + uint32_t head_dim_qk = q.size(2); + uint32_t head_dim_vo = paged_v_cache.size(3); + + TORCH_CHECK(head_dim_qk == head_dim_vo, + "CUDA cores template only supports equal head dim for QK and VO, please use tensor " + "cores template for different head dim"); + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == batch_size, lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == num_qo_heads, lse.size(1), q.size(1)); + } + + void* float_buffer = static_cast(float_workspace_buffer.data_ptr()); + void* int_buffer = static_cast(int_workspace_buffer.data_ptr()); + + // get q_scalar_type and kv_scalar_type + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = paged_k_cache.scalar_type(); + + // get q_stride_n and q_stride_h + const auto q_stride_n = q.stride(0); + const auto q_stride_h = q.stride(1); + + // get kv_cache_strides + const int64_t* kv_cache_strides = nullptr; + auto k_strides = paged_k_cache.strides(); + auto v_strides = paged_v_cache.strides(); + TORCH_CHECK(k_strides == v_strides, "k/v strides must be identical"); + kv_cache_strides = k_strides.data(); + + const c10::cuda::OptionalCUDAGuard device_guard(device); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + paged_kv_t paged_kv( + num_kv_heads, page_size, HEAD_DIM_QK, batch_size, kv_layout, + static_cast(paged_k_cache.data_ptr()), + static_cast(paged_v_cache.data_ptr()), kv_cache_strides, + static_cast(paged_kv_indices.data_ptr()), + static_cast(paged_kv_indptr.data_ptr()), + static_cast(paged_kv_last_page_len.data_ptr())); + + Params params; + params.q = static_cast(q.data_ptr()); + params.paged_kv = paged_kv; + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.padded_batch_size = 0; + params.num_qo_heads = num_qo_heads; + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.window_left = window_left; + params.request_indices = nullptr; + params.kv_tile_indices = nullptr; + params.o_indptr = nullptr; + params.kv_chunk_size_ptr = nullptr; + params.block_valid_mask = nullptr; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + DTypeO* tmp_v = nullptr; + float* tmp_s = nullptr; + params.request_indices = + GetPtrFromBaseOffset(int_buffer, plan_info.request_indices_offset); + params.kv_tile_indices = + GetPtrFromBaseOffset(int_buffer, plan_info.kv_tile_indices_offset); + params.o_indptr = GetPtrFromBaseOffset(int_buffer, plan_info.o_indptr_offset); + params.kv_chunk_size_ptr = + GetPtrFromBaseOffset(int_buffer, plan_info.kv_chunk_size_ptr_offset); + if (plan_info.split_kv) { + tmp_v = GetPtrFromBaseOffset(float_buffer, plan_info.v_offset); + tmp_s = GetPtrFromBaseOffset(float_buffer, plan_info.s_offset); + if (plan_info.enable_cuda_graph) { + params.block_valid_mask = + GetPtrFromBaseOffset(int_buffer, plan_info.block_valid_mask_offset); + } + } + params.padded_batch_size = plan_info.padded_batch_size; + + cudaError_t status = + flashinfer::BatchDecodeWithPagedKVCacheDispatched(params, tmp_v, + tmp_s, enable_pdl, + /*stream=*/stream); + TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); +} diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_config.inc b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..6b43ef26d8c6d6ec4d821a2eb7744d0ce6549273 --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_config.inc @@ -0,0 +1,71 @@ +#pragma once +#include +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta +#define ADDITIONAL_PARAMS_SETTER params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, ...) { \ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ +} + +using namespace flashinfer; + +using DTypeQ = nv_bfloat16; +using DTypeKV = nv_bfloat16; +using DTypeO = nv_bfloat16; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 64; +constexpr int HEAD_DIM_VO = 64; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = false; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + DTypeQ* q; + paged_kv_t paged_kv; + DTypeO* o; + float* lse; + + float* maybe_alibi_slopes; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; + + + uint32_t padded_batch_size; + uint32_t num_qo_heads; + IdType q_stride_n; + IdType q_stride_h; + int32_t window_left; + bool enable_pdl; + + IdType* request_indices; + IdType* kv_tile_indices; + IdType* o_indptr; + IdType* kv_chunk_size_ptr; + bool* block_valid_mask; + bool partition_kv; + + __host__ __device__ __forceinline__ int32_t get_qo_len(int32_t batch_idx) const { return 1; } + + __host__ __device__ __forceinline__ int32_t get_kv_len(int32_t batch_idx) const { + return paged_kv.get_length(batch_idx); + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_jit_pybind.cu b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..91e944770b8cee57481f0631545cef02694f50cc --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_jit_pybind.cu @@ -0,0 +1,40 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "batch_decode_config.inc" +#include "pytorch_extension_utils.h" + +at::Tensor BatchDecodeWithPagedKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size, + int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, + int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo, + at::Tensor empty_q_data, at::Tensor empty_kv_data); + +void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, + at::Tensor paged_v_cache, at::Tensor paged_kv_indptr, + at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, + int64_t kv_layout_code, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Batched decode with paged KV-Cache plan + m.def("plan", BatchDecodeWithPagedKVCachePlan); + // Batched decode with paged KV-Cache run + m.def("run", BatchDecodeWithPagedKVCacheRun); +} diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_kernel.cu b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_kernel.cu new file mode 100644 index 0000000000000000000000000000000000000000..9ba3df689075c5c15e90cb5ff527d8883eeef296 --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_kernel.cu @@ -0,0 +1,13 @@ +#include +#include "batch_decode_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t +BatchDecodeWithPagedKVCacheDispatched<64, PosEncodingMode::kNone, DefaultAttention, Params>( + Params params, nv_bfloat16* tmp_v, + float* tmp_s, bool enable_pdl, cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode.cu b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode.cu new file mode 100644 index 0000000000000000000000000000000000000000..02c75e02718f8f8ce9b6036f8f788d9a9e254cf1 --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode.cu @@ -0,0 +1,197 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include +#include + +#include "batch_decode_config.inc" +#include "pytorch_conversion_utils.h" +#include "pytorch_extension_utils.h" + +namespace flashinfer { + +template +cudaError_t BatchDecodeWithPagedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v, + float* tmp_s, bool enable_pdl, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +at::Tensor BatchDecodeWithPagedKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size, + int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, + int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo, + at::Tensor empty_q_data, at::Tensor empty_kv_data) { + size_t float_workspace_size_in_bytes = + float_workspace_buffer.size(0) * float_workspace_buffer.element_size(); + size_t int_workspace_size_in_bytes = + int_workspace_buffer.size(0) * int_workspace_buffer.element_size(); + + DecodePlanInfo plan_info; + + auto q_scalar_type = empty_q_data.scalar_type(); + auto kv_scalar_type = empty_kv_data.scalar_type(); + + TORCH_CHECK(head_dim_qk == head_dim_vo, + "CUDA cores template only supports equal head dim for QK and VO, please use tensor " + "cores template for different head dim"); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + DISPATCH_GQA_GROUP_SIZE(num_qo_heads / num_kv_heads, GROUP_SIZE, { + auto work_estimation_func = BatchDecodeWithPagedKVCacheWorkEstimationDispatched< + GROUP_SIZE, HEAD_DIM_QK, POS_ENCODING_MODE, AttentionVariant, Params>; + cudaError_t status = DecodePlan( + static_cast(float_workspace_buffer.data_ptr()), float_workspace_size_in_bytes, + static_cast(int_workspace_buffer.data_ptr()), + static_cast(page_locked_int_workspace_buffer.data_ptr()), + int_workspace_size_in_bytes, plan_info, static_cast(indptr.data_ptr()), + batch_size, num_qo_heads, page_size, enable_cuda_graph, + /*stream=*/stream, work_estimation_func); + + TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); + }); + + return vec_to_tensor(plan_info.ToVector()); +} + +void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, + at::Tensor paged_v_cache, at::Tensor paged_kv_indptr, + at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, + int64_t kv_layout_code, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + DecodePlanInfo plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + QKVLayout kv_layout = static_cast(kv_layout_code); + auto device = q.device(); + int64_t batch_size = q.size(0); + int64_t num_qo_heads = q.size(1); + int64_t num_kv_heads, page_size; + + if (kv_layout == QKVLayout::kHND) { + num_kv_heads = paged_k_cache.size(1); + page_size = paged_k_cache.size(2); + } else { + page_size = paged_k_cache.size(1); + num_kv_heads = paged_k_cache.size(2); + } + uint32_t head_dim_qk = q.size(2); + uint32_t head_dim_vo = paged_v_cache.size(3); + + TORCH_CHECK(head_dim_qk == head_dim_vo, + "CUDA cores template only supports equal head dim for QK and VO, please use tensor " + "cores template for different head dim"); + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == batch_size, lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == num_qo_heads, lse.size(1), q.size(1)); + } + + void* float_buffer = static_cast(float_workspace_buffer.data_ptr()); + void* int_buffer = static_cast(int_workspace_buffer.data_ptr()); + + // get q_scalar_type and kv_scalar_type + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = paged_k_cache.scalar_type(); + + // get q_stride_n and q_stride_h + const auto q_stride_n = q.stride(0); + const auto q_stride_h = q.stride(1); + + // get kv_cache_strides + const int64_t* kv_cache_strides = nullptr; + auto k_strides = paged_k_cache.strides(); + auto v_strides = paged_v_cache.strides(); + TORCH_CHECK(k_strides == v_strides, "k/v strides must be identical"); + kv_cache_strides = k_strides.data(); + + const c10::cuda::OptionalCUDAGuard device_guard(device); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + paged_kv_t paged_kv( + num_kv_heads, page_size, HEAD_DIM_QK, batch_size, kv_layout, + static_cast(paged_k_cache.data_ptr()), + static_cast(paged_v_cache.data_ptr()), kv_cache_strides, + static_cast(paged_kv_indices.data_ptr()), + static_cast(paged_kv_indptr.data_ptr()), + static_cast(paged_kv_last_page_len.data_ptr())); + + Params params; + params.q = static_cast(q.data_ptr()); + params.paged_kv = paged_kv; + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.padded_batch_size = 0; + params.num_qo_heads = num_qo_heads; + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.window_left = window_left; + params.request_indices = nullptr; + params.kv_tile_indices = nullptr; + params.o_indptr = nullptr; + params.kv_chunk_size_ptr = nullptr; + params.block_valid_mask = nullptr; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + DTypeO* tmp_v = nullptr; + float* tmp_s = nullptr; + params.request_indices = + GetPtrFromBaseOffset(int_buffer, plan_info.request_indices_offset); + params.kv_tile_indices = + GetPtrFromBaseOffset(int_buffer, plan_info.kv_tile_indices_offset); + params.o_indptr = GetPtrFromBaseOffset(int_buffer, plan_info.o_indptr_offset); + params.kv_chunk_size_ptr = + GetPtrFromBaseOffset(int_buffer, plan_info.kv_chunk_size_ptr_offset); + if (plan_info.split_kv) { + tmp_v = GetPtrFromBaseOffset(float_buffer, plan_info.v_offset); + tmp_s = GetPtrFromBaseOffset(float_buffer, plan_info.s_offset); + if (plan_info.enable_cuda_graph) { + params.block_valid_mask = + GetPtrFromBaseOffset(int_buffer, plan_info.block_valid_mask_offset); + } + } + params.padded_batch_size = plan_info.padded_batch_size; + + cudaError_t status = + flashinfer::BatchDecodeWithPagedKVCacheDispatched(params, tmp_v, + tmp_s, enable_pdl, + /*stream=*/stream); + TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); +} diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode_config.inc b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..ae3692fd9aeb200b226917f8c8e4048a6d0c7d9d --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode_config.inc @@ -0,0 +1,71 @@ +#pragma once +#include +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta +#define ADDITIONAL_PARAMS_SETTER params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, ...) { \ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ +} + +using namespace flashinfer; + +using DTypeQ = nv_bfloat16; +using DTypeKV = nv_bfloat16; +using DTypeO = nv_bfloat16; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 64; +constexpr int HEAD_DIM_VO = 64; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = true; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + DTypeQ* q; + paged_kv_t paged_kv; + DTypeO* o; + float* lse; + + float* maybe_alibi_slopes; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; + + + uint32_t padded_batch_size; + uint32_t num_qo_heads; + IdType q_stride_n; + IdType q_stride_h; + int32_t window_left; + bool enable_pdl; + + IdType* request_indices; + IdType* kv_tile_indices; + IdType* o_indptr; + IdType* kv_chunk_size_ptr; + bool* block_valid_mask; + bool partition_kv; + + __host__ __device__ __forceinline__ int32_t get_qo_len(int32_t batch_idx) const { return 1; } + + __host__ __device__ __forceinline__ int32_t get_kv_len(int32_t batch_idx) const { + return paged_kv.get_length(batch_idx); + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode_jit_pybind.cu b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..91e944770b8cee57481f0631545cef02694f50cc --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode_jit_pybind.cu @@ -0,0 +1,40 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "batch_decode_config.inc" +#include "pytorch_extension_utils.h" + +at::Tensor BatchDecodeWithPagedKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size, + int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, + int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo, + at::Tensor empty_q_data, at::Tensor empty_kv_data); + +void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, + at::Tensor paged_v_cache, at::Tensor paged_kv_indptr, + at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, + int64_t kv_layout_code, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Batched decode with paged KV-Cache plan + m.def("plan", BatchDecodeWithPagedKVCachePlan); + // Batched decode with paged KV-Cache run + m.def("run", BatchDecodeWithPagedKVCacheRun); +} diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode_kernel.cu b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode_kernel.cu new file mode 100644 index 0000000000000000000000000000000000000000..6cb7d97320a918556f4e46afbaba5444a0334225 --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode_kernel.cu @@ -0,0 +1,13 @@ +#include +#include "batch_decode_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t +BatchDecodeWithPagedKVCacheDispatched<64, PosEncodingMode::kNone, DefaultAttention, Params>( + Params params, nv_bfloat16* tmp_v, + float* tmp_s, bool enable_pdl, cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode.cu b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode.cu new file mode 100644 index 0000000000000000000000000000000000000000..02c75e02718f8f8ce9b6036f8f788d9a9e254cf1 --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode.cu @@ -0,0 +1,197 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include +#include + +#include "batch_decode_config.inc" +#include "pytorch_conversion_utils.h" +#include "pytorch_extension_utils.h" + +namespace flashinfer { + +template +cudaError_t BatchDecodeWithPagedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v, + float* tmp_s, bool enable_pdl, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +at::Tensor BatchDecodeWithPagedKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size, + int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, + int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo, + at::Tensor empty_q_data, at::Tensor empty_kv_data) { + size_t float_workspace_size_in_bytes = + float_workspace_buffer.size(0) * float_workspace_buffer.element_size(); + size_t int_workspace_size_in_bytes = + int_workspace_buffer.size(0) * int_workspace_buffer.element_size(); + + DecodePlanInfo plan_info; + + auto q_scalar_type = empty_q_data.scalar_type(); + auto kv_scalar_type = empty_kv_data.scalar_type(); + + TORCH_CHECK(head_dim_qk == head_dim_vo, + "CUDA cores template only supports equal head dim for QK and VO, please use tensor " + "cores template for different head dim"); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + DISPATCH_GQA_GROUP_SIZE(num_qo_heads / num_kv_heads, GROUP_SIZE, { + auto work_estimation_func = BatchDecodeWithPagedKVCacheWorkEstimationDispatched< + GROUP_SIZE, HEAD_DIM_QK, POS_ENCODING_MODE, AttentionVariant, Params>; + cudaError_t status = DecodePlan( + static_cast(float_workspace_buffer.data_ptr()), float_workspace_size_in_bytes, + static_cast(int_workspace_buffer.data_ptr()), + static_cast(page_locked_int_workspace_buffer.data_ptr()), + int_workspace_size_in_bytes, plan_info, static_cast(indptr.data_ptr()), + batch_size, num_qo_heads, page_size, enable_cuda_graph, + /*stream=*/stream, work_estimation_func); + + TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); + }); + + return vec_to_tensor(plan_info.ToVector()); +} + +void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, + at::Tensor paged_v_cache, at::Tensor paged_kv_indptr, + at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, + int64_t kv_layout_code, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + DecodePlanInfo plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + QKVLayout kv_layout = static_cast(kv_layout_code); + auto device = q.device(); + int64_t batch_size = q.size(0); + int64_t num_qo_heads = q.size(1); + int64_t num_kv_heads, page_size; + + if (kv_layout == QKVLayout::kHND) { + num_kv_heads = paged_k_cache.size(1); + page_size = paged_k_cache.size(2); + } else { + page_size = paged_k_cache.size(1); + num_kv_heads = paged_k_cache.size(2); + } + uint32_t head_dim_qk = q.size(2); + uint32_t head_dim_vo = paged_v_cache.size(3); + + TORCH_CHECK(head_dim_qk == head_dim_vo, + "CUDA cores template only supports equal head dim for QK and VO, please use tensor " + "cores template for different head dim"); + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == batch_size, lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == num_qo_heads, lse.size(1), q.size(1)); + } + + void* float_buffer = static_cast(float_workspace_buffer.data_ptr()); + void* int_buffer = static_cast(int_workspace_buffer.data_ptr()); + + // get q_scalar_type and kv_scalar_type + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = paged_k_cache.scalar_type(); + + // get q_stride_n and q_stride_h + const auto q_stride_n = q.stride(0); + const auto q_stride_h = q.stride(1); + + // get kv_cache_strides + const int64_t* kv_cache_strides = nullptr; + auto k_strides = paged_k_cache.strides(); + auto v_strides = paged_v_cache.strides(); + TORCH_CHECK(k_strides == v_strides, "k/v strides must be identical"); + kv_cache_strides = k_strides.data(); + + const c10::cuda::OptionalCUDAGuard device_guard(device); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + paged_kv_t paged_kv( + num_kv_heads, page_size, HEAD_DIM_QK, batch_size, kv_layout, + static_cast(paged_k_cache.data_ptr()), + static_cast(paged_v_cache.data_ptr()), kv_cache_strides, + static_cast(paged_kv_indices.data_ptr()), + static_cast(paged_kv_indptr.data_ptr()), + static_cast(paged_kv_last_page_len.data_ptr())); + + Params params; + params.q = static_cast(q.data_ptr()); + params.paged_kv = paged_kv; + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.padded_batch_size = 0; + params.num_qo_heads = num_qo_heads; + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.window_left = window_left; + params.request_indices = nullptr; + params.kv_tile_indices = nullptr; + params.o_indptr = nullptr; + params.kv_chunk_size_ptr = nullptr; + params.block_valid_mask = nullptr; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + DTypeO* tmp_v = nullptr; + float* tmp_s = nullptr; + params.request_indices = + GetPtrFromBaseOffset(int_buffer, plan_info.request_indices_offset); + params.kv_tile_indices = + GetPtrFromBaseOffset(int_buffer, plan_info.kv_tile_indices_offset); + params.o_indptr = GetPtrFromBaseOffset(int_buffer, plan_info.o_indptr_offset); + params.kv_chunk_size_ptr = + GetPtrFromBaseOffset(int_buffer, plan_info.kv_chunk_size_ptr_offset); + if (plan_info.split_kv) { + tmp_v = GetPtrFromBaseOffset(float_buffer, plan_info.v_offset); + tmp_s = GetPtrFromBaseOffset(float_buffer, plan_info.s_offset); + if (plan_info.enable_cuda_graph) { + params.block_valid_mask = + GetPtrFromBaseOffset(int_buffer, plan_info.block_valid_mask_offset); + } + } + params.padded_batch_size = plan_info.padded_batch_size; + + cudaError_t status = + flashinfer::BatchDecodeWithPagedKVCacheDispatched(params, tmp_v, + tmp_s, enable_pdl, + /*stream=*/stream); + TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); +} diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_config.inc b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..b0d70466bce7b48b0d2e0beb6128aea5a05e1795 --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_config.inc @@ -0,0 +1,71 @@ +#pragma once +#include +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta +#define ADDITIONAL_PARAMS_SETTER params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, ...) { \ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ +} + +using namespace flashinfer; + +using DTypeQ = nv_bfloat16; +using DTypeKV = __nv_fp8_e4m3; +using DTypeO = nv_bfloat16; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 128; +constexpr int HEAD_DIM_VO = 128; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = false; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + DTypeQ* q; + paged_kv_t paged_kv; + DTypeO* o; + float* lse; + + float* maybe_alibi_slopes; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; + + + uint32_t padded_batch_size; + uint32_t num_qo_heads; + IdType q_stride_n; + IdType q_stride_h; + int32_t window_left; + bool enable_pdl; + + IdType* request_indices; + IdType* kv_tile_indices; + IdType* o_indptr; + IdType* kv_chunk_size_ptr; + bool* block_valid_mask; + bool partition_kv; + + __host__ __device__ __forceinline__ int32_t get_qo_len(int32_t batch_idx) const { return 1; } + + __host__ __device__ __forceinline__ int32_t get_kv_len(int32_t batch_idx) const { + return paged_kv.get_length(batch_idx); + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_jit_pybind.cu b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..91e944770b8cee57481f0631545cef02694f50cc --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_jit_pybind.cu @@ -0,0 +1,40 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "batch_decode_config.inc" +#include "pytorch_extension_utils.h" + +at::Tensor BatchDecodeWithPagedKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size, + int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, + int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo, + at::Tensor empty_q_data, at::Tensor empty_kv_data); + +void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, + at::Tensor paged_v_cache, at::Tensor paged_kv_indptr, + at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, + int64_t kv_layout_code, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Batched decode with paged KV-Cache plan + m.def("plan", BatchDecodeWithPagedKVCachePlan); + // Batched decode with paged KV-Cache run + m.def("run", BatchDecodeWithPagedKVCacheRun); +} diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_kernel.cu b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_kernel.cu new file mode 100644 index 0000000000000000000000000000000000000000..973142e154c74a17ca0e21d47d510e1167de4891 --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_kernel.cu @@ -0,0 +1,13 @@ +#include +#include "batch_decode_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t +BatchDecodeWithPagedKVCacheDispatched<128, PosEncodingMode::kNone, DefaultAttention, Params>( + Params params, nv_bfloat16* tmp_v, + float* tmp_s, bool enable_pdl, cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode.cu b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode.cu new file mode 100644 index 0000000000000000000000000000000000000000..02c75e02718f8f8ce9b6036f8f788d9a9e254cf1 --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode.cu @@ -0,0 +1,197 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include +#include + +#include "batch_decode_config.inc" +#include "pytorch_conversion_utils.h" +#include "pytorch_extension_utils.h" + +namespace flashinfer { + +template +cudaError_t BatchDecodeWithPagedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v, + float* tmp_s, bool enable_pdl, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +at::Tensor BatchDecodeWithPagedKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size, + int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, + int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo, + at::Tensor empty_q_data, at::Tensor empty_kv_data) { + size_t float_workspace_size_in_bytes = + float_workspace_buffer.size(0) * float_workspace_buffer.element_size(); + size_t int_workspace_size_in_bytes = + int_workspace_buffer.size(0) * int_workspace_buffer.element_size(); + + DecodePlanInfo plan_info; + + auto q_scalar_type = empty_q_data.scalar_type(); + auto kv_scalar_type = empty_kv_data.scalar_type(); + + TORCH_CHECK(head_dim_qk == head_dim_vo, + "CUDA cores template only supports equal head dim for QK and VO, please use tensor " + "cores template for different head dim"); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + DISPATCH_GQA_GROUP_SIZE(num_qo_heads / num_kv_heads, GROUP_SIZE, { + auto work_estimation_func = BatchDecodeWithPagedKVCacheWorkEstimationDispatched< + GROUP_SIZE, HEAD_DIM_QK, POS_ENCODING_MODE, AttentionVariant, Params>; + cudaError_t status = DecodePlan( + static_cast(float_workspace_buffer.data_ptr()), float_workspace_size_in_bytes, + static_cast(int_workspace_buffer.data_ptr()), + static_cast(page_locked_int_workspace_buffer.data_ptr()), + int_workspace_size_in_bytes, plan_info, static_cast(indptr.data_ptr()), + batch_size, num_qo_heads, page_size, enable_cuda_graph, + /*stream=*/stream, work_estimation_func); + + TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); + }); + + return vec_to_tensor(plan_info.ToVector()); +} + +void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, + at::Tensor paged_v_cache, at::Tensor paged_kv_indptr, + at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, + int64_t kv_layout_code, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + DecodePlanInfo plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + QKVLayout kv_layout = static_cast(kv_layout_code); + auto device = q.device(); + int64_t batch_size = q.size(0); + int64_t num_qo_heads = q.size(1); + int64_t num_kv_heads, page_size; + + if (kv_layout == QKVLayout::kHND) { + num_kv_heads = paged_k_cache.size(1); + page_size = paged_k_cache.size(2); + } else { + page_size = paged_k_cache.size(1); + num_kv_heads = paged_k_cache.size(2); + } + uint32_t head_dim_qk = q.size(2); + uint32_t head_dim_vo = paged_v_cache.size(3); + + TORCH_CHECK(head_dim_qk == head_dim_vo, + "CUDA cores template only supports equal head dim for QK and VO, please use tensor " + "cores template for different head dim"); + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == batch_size, lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == num_qo_heads, lse.size(1), q.size(1)); + } + + void* float_buffer = static_cast(float_workspace_buffer.data_ptr()); + void* int_buffer = static_cast(int_workspace_buffer.data_ptr()); + + // get q_scalar_type and kv_scalar_type + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = paged_k_cache.scalar_type(); + + // get q_stride_n and q_stride_h + const auto q_stride_n = q.stride(0); + const auto q_stride_h = q.stride(1); + + // get kv_cache_strides + const int64_t* kv_cache_strides = nullptr; + auto k_strides = paged_k_cache.strides(); + auto v_strides = paged_v_cache.strides(); + TORCH_CHECK(k_strides == v_strides, "k/v strides must be identical"); + kv_cache_strides = k_strides.data(); + + const c10::cuda::OptionalCUDAGuard device_guard(device); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + paged_kv_t paged_kv( + num_kv_heads, page_size, HEAD_DIM_QK, batch_size, kv_layout, + static_cast(paged_k_cache.data_ptr()), + static_cast(paged_v_cache.data_ptr()), kv_cache_strides, + static_cast(paged_kv_indices.data_ptr()), + static_cast(paged_kv_indptr.data_ptr()), + static_cast(paged_kv_last_page_len.data_ptr())); + + Params params; + params.q = static_cast(q.data_ptr()); + params.paged_kv = paged_kv; + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.padded_batch_size = 0; + params.num_qo_heads = num_qo_heads; + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.window_left = window_left; + params.request_indices = nullptr; + params.kv_tile_indices = nullptr; + params.o_indptr = nullptr; + params.kv_chunk_size_ptr = nullptr; + params.block_valid_mask = nullptr; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + DTypeO* tmp_v = nullptr; + float* tmp_s = nullptr; + params.request_indices = + GetPtrFromBaseOffset(int_buffer, plan_info.request_indices_offset); + params.kv_tile_indices = + GetPtrFromBaseOffset(int_buffer, plan_info.kv_tile_indices_offset); + params.o_indptr = GetPtrFromBaseOffset(int_buffer, plan_info.o_indptr_offset); + params.kv_chunk_size_ptr = + GetPtrFromBaseOffset(int_buffer, plan_info.kv_chunk_size_ptr_offset); + if (plan_info.split_kv) { + tmp_v = GetPtrFromBaseOffset(float_buffer, plan_info.v_offset); + tmp_s = GetPtrFromBaseOffset(float_buffer, plan_info.s_offset); + if (plan_info.enable_cuda_graph) { + params.block_valid_mask = + GetPtrFromBaseOffset(int_buffer, plan_info.block_valid_mask_offset); + } + } + params.padded_batch_size = plan_info.padded_batch_size; + + cudaError_t status = + flashinfer::BatchDecodeWithPagedKVCacheDispatched(params, tmp_v, + tmp_s, enable_pdl, + /*stream=*/stream); + TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); +} diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode_config.inc b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..a5c653eeb45c68319c3f542aa31a15db0003f5ae --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode_config.inc @@ -0,0 +1,71 @@ +#pragma once +#include +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta +#define ADDITIONAL_PARAMS_SETTER params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, ...) { \ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ +} + +using namespace flashinfer; + +using DTypeQ = nv_bfloat16; +using DTypeKV = __nv_fp8_e4m3; +using DTypeO = nv_bfloat16; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 256; +constexpr int HEAD_DIM_VO = 256; +constexpr auto USE_LOGITS_SOFT_CAP = true; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = true; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + DTypeQ* q; + paged_kv_t paged_kv; + DTypeO* o; + float* lse; + + float* maybe_alibi_slopes; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; + + + uint32_t padded_batch_size; + uint32_t num_qo_heads; + IdType q_stride_n; + IdType q_stride_h; + int32_t window_left; + bool enable_pdl; + + IdType* request_indices; + IdType* kv_tile_indices; + IdType* o_indptr; + IdType* kv_chunk_size_ptr; + bool* block_valid_mask; + bool partition_kv; + + __host__ __device__ __forceinline__ int32_t get_qo_len(int32_t batch_idx) const { return 1; } + + __host__ __device__ __forceinline__ int32_t get_kv_len(int32_t batch_idx) const { + return paged_kv.get_length(batch_idx); + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode_jit_pybind.cu b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..91e944770b8cee57481f0631545cef02694f50cc --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode_jit_pybind.cu @@ -0,0 +1,40 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "batch_decode_config.inc" +#include "pytorch_extension_utils.h" + +at::Tensor BatchDecodeWithPagedKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size, + int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, + int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo, + at::Tensor empty_q_data, at::Tensor empty_kv_data); + +void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, + at::Tensor paged_v_cache, at::Tensor paged_kv_indptr, + at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, + int64_t kv_layout_code, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Batched decode with paged KV-Cache plan + m.def("plan", BatchDecodeWithPagedKVCachePlan); + // Batched decode with paged KV-Cache run + m.def("run", BatchDecodeWithPagedKVCacheRun); +} diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode_kernel.cu b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode_kernel.cu new file mode 100644 index 0000000000000000000000000000000000000000..b63da6a66b8eeb551357ef21128a476cd6477474 --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode_kernel.cu @@ -0,0 +1,13 @@ +#include +#include "batch_decode_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t +BatchDecodeWithPagedKVCacheDispatched<256, PosEncodingMode::kNone, DefaultAttention, Params>( + Params params, nv_bfloat16* tmp_v, + float* tmp_s, bool enable_pdl, cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode.cu b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode.cu new file mode 100644 index 0000000000000000000000000000000000000000..02c75e02718f8f8ce9b6036f8f788d9a9e254cf1 --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode.cu @@ -0,0 +1,197 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include +#include + +#include "batch_decode_config.inc" +#include "pytorch_conversion_utils.h" +#include "pytorch_extension_utils.h" + +namespace flashinfer { + +template +cudaError_t BatchDecodeWithPagedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v, + float* tmp_s, bool enable_pdl, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +at::Tensor BatchDecodeWithPagedKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size, + int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, + int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo, + at::Tensor empty_q_data, at::Tensor empty_kv_data) { + size_t float_workspace_size_in_bytes = + float_workspace_buffer.size(0) * float_workspace_buffer.element_size(); + size_t int_workspace_size_in_bytes = + int_workspace_buffer.size(0) * int_workspace_buffer.element_size(); + + DecodePlanInfo plan_info; + + auto q_scalar_type = empty_q_data.scalar_type(); + auto kv_scalar_type = empty_kv_data.scalar_type(); + + TORCH_CHECK(head_dim_qk == head_dim_vo, + "CUDA cores template only supports equal head dim for QK and VO, please use tensor " + "cores template for different head dim"); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + DISPATCH_GQA_GROUP_SIZE(num_qo_heads / num_kv_heads, GROUP_SIZE, { + auto work_estimation_func = BatchDecodeWithPagedKVCacheWorkEstimationDispatched< + GROUP_SIZE, HEAD_DIM_QK, POS_ENCODING_MODE, AttentionVariant, Params>; + cudaError_t status = DecodePlan( + static_cast(float_workspace_buffer.data_ptr()), float_workspace_size_in_bytes, + static_cast(int_workspace_buffer.data_ptr()), + static_cast(page_locked_int_workspace_buffer.data_ptr()), + int_workspace_size_in_bytes, plan_info, static_cast(indptr.data_ptr()), + batch_size, num_qo_heads, page_size, enable_cuda_graph, + /*stream=*/stream, work_estimation_func); + + TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); + }); + + return vec_to_tensor(plan_info.ToVector()); +} + +void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, + at::Tensor paged_v_cache, at::Tensor paged_kv_indptr, + at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, + int64_t kv_layout_code, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + DecodePlanInfo plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + QKVLayout kv_layout = static_cast(kv_layout_code); + auto device = q.device(); + int64_t batch_size = q.size(0); + int64_t num_qo_heads = q.size(1); + int64_t num_kv_heads, page_size; + + if (kv_layout == QKVLayout::kHND) { + num_kv_heads = paged_k_cache.size(1); + page_size = paged_k_cache.size(2); + } else { + page_size = paged_k_cache.size(1); + num_kv_heads = paged_k_cache.size(2); + } + uint32_t head_dim_qk = q.size(2); + uint32_t head_dim_vo = paged_v_cache.size(3); + + TORCH_CHECK(head_dim_qk == head_dim_vo, + "CUDA cores template only supports equal head dim for QK and VO, please use tensor " + "cores template for different head dim"); + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == batch_size, lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == num_qo_heads, lse.size(1), q.size(1)); + } + + void* float_buffer = static_cast(float_workspace_buffer.data_ptr()); + void* int_buffer = static_cast(int_workspace_buffer.data_ptr()); + + // get q_scalar_type and kv_scalar_type + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = paged_k_cache.scalar_type(); + + // get q_stride_n and q_stride_h + const auto q_stride_n = q.stride(0); + const auto q_stride_h = q.stride(1); + + // get kv_cache_strides + const int64_t* kv_cache_strides = nullptr; + auto k_strides = paged_k_cache.strides(); + auto v_strides = paged_v_cache.strides(); + TORCH_CHECK(k_strides == v_strides, "k/v strides must be identical"); + kv_cache_strides = k_strides.data(); + + const c10::cuda::OptionalCUDAGuard device_guard(device); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + paged_kv_t paged_kv( + num_kv_heads, page_size, HEAD_DIM_QK, batch_size, kv_layout, + static_cast(paged_k_cache.data_ptr()), + static_cast(paged_v_cache.data_ptr()), kv_cache_strides, + static_cast(paged_kv_indices.data_ptr()), + static_cast(paged_kv_indptr.data_ptr()), + static_cast(paged_kv_last_page_len.data_ptr())); + + Params params; + params.q = static_cast(q.data_ptr()); + params.paged_kv = paged_kv; + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.padded_batch_size = 0; + params.num_qo_heads = num_qo_heads; + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.window_left = window_left; + params.request_indices = nullptr; + params.kv_tile_indices = nullptr; + params.o_indptr = nullptr; + params.kv_chunk_size_ptr = nullptr; + params.block_valid_mask = nullptr; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + DTypeO* tmp_v = nullptr; + float* tmp_s = nullptr; + params.request_indices = + GetPtrFromBaseOffset(int_buffer, plan_info.request_indices_offset); + params.kv_tile_indices = + GetPtrFromBaseOffset(int_buffer, plan_info.kv_tile_indices_offset); + params.o_indptr = GetPtrFromBaseOffset(int_buffer, plan_info.o_indptr_offset); + params.kv_chunk_size_ptr = + GetPtrFromBaseOffset(int_buffer, plan_info.kv_chunk_size_ptr_offset); + if (plan_info.split_kv) { + tmp_v = GetPtrFromBaseOffset(float_buffer, plan_info.v_offset); + tmp_s = GetPtrFromBaseOffset(float_buffer, plan_info.s_offset); + if (plan_info.enable_cuda_graph) { + params.block_valid_mask = + GetPtrFromBaseOffset(int_buffer, plan_info.block_valid_mask_offset); + } + } + params.padded_batch_size = plan_info.padded_batch_size; + + cudaError_t status = + flashinfer::BatchDecodeWithPagedKVCacheDispatched(params, tmp_v, + tmp_s, enable_pdl, + /*stream=*/stream); + TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); +} diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_config.inc b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..8fc471fdb4b28e126316be9ec730207f2d149eb1 --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_config.inc @@ -0,0 +1,71 @@ +#pragma once +#include +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta +#define ADDITIONAL_PARAMS_SETTER params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, ...) { \ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ +} + +using namespace flashinfer; + +using DTypeQ = nv_bfloat16; +using DTypeKV = __nv_fp8_e4m3; +using DTypeO = nv_bfloat16; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 64; +constexpr int HEAD_DIM_VO = 64; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = false; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + DTypeQ* q; + paged_kv_t paged_kv; + DTypeO* o; + float* lse; + + float* maybe_alibi_slopes; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; + + + uint32_t padded_batch_size; + uint32_t num_qo_heads; + IdType q_stride_n; + IdType q_stride_h; + int32_t window_left; + bool enable_pdl; + + IdType* request_indices; + IdType* kv_tile_indices; + IdType* o_indptr; + IdType* kv_chunk_size_ptr; + bool* block_valid_mask; + bool partition_kv; + + __host__ __device__ __forceinline__ int32_t get_qo_len(int32_t batch_idx) const { return 1; } + + __host__ __device__ __forceinline__ int32_t get_kv_len(int32_t batch_idx) const { + return paged_kv.get_length(batch_idx); + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_jit_pybind.cu b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..91e944770b8cee57481f0631545cef02694f50cc --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_jit_pybind.cu @@ -0,0 +1,40 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "batch_decode_config.inc" +#include "pytorch_extension_utils.h" + +at::Tensor BatchDecodeWithPagedKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size, + int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, + int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo, + at::Tensor empty_q_data, at::Tensor empty_kv_data); + +void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, + at::Tensor paged_v_cache, at::Tensor paged_kv_indptr, + at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, + int64_t kv_layout_code, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Batched decode with paged KV-Cache plan + m.def("plan", BatchDecodeWithPagedKVCachePlan); + // Batched decode with paged KV-Cache run + m.def("run", BatchDecodeWithPagedKVCacheRun); +} diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_kernel.cu b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_kernel.cu new file mode 100644 index 0000000000000000000000000000000000000000..9ba3df689075c5c15e90cb5ff527d8883eeef296 --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_kernel.cu @@ -0,0 +1,13 @@ +#include +#include "batch_decode_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t +BatchDecodeWithPagedKVCacheDispatched<64, PosEncodingMode::kNone, DefaultAttention, Params>( + Params params, nv_bfloat16* tmp_v, + float* tmp_s, bool enable_pdl, cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode.cu b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode.cu new file mode 100644 index 0000000000000000000000000000000000000000..02c75e02718f8f8ce9b6036f8f788d9a9e254cf1 --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode.cu @@ -0,0 +1,197 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include +#include + +#include "batch_decode_config.inc" +#include "pytorch_conversion_utils.h" +#include "pytorch_extension_utils.h" + +namespace flashinfer { + +template +cudaError_t BatchDecodeWithPagedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v, + float* tmp_s, bool enable_pdl, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +at::Tensor BatchDecodeWithPagedKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size, + int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, + int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo, + at::Tensor empty_q_data, at::Tensor empty_kv_data) { + size_t float_workspace_size_in_bytes = + float_workspace_buffer.size(0) * float_workspace_buffer.element_size(); + size_t int_workspace_size_in_bytes = + int_workspace_buffer.size(0) * int_workspace_buffer.element_size(); + + DecodePlanInfo plan_info; + + auto q_scalar_type = empty_q_data.scalar_type(); + auto kv_scalar_type = empty_kv_data.scalar_type(); + + TORCH_CHECK(head_dim_qk == head_dim_vo, + "CUDA cores template only supports equal head dim for QK and VO, please use tensor " + "cores template for different head dim"); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + DISPATCH_GQA_GROUP_SIZE(num_qo_heads / num_kv_heads, GROUP_SIZE, { + auto work_estimation_func = BatchDecodeWithPagedKVCacheWorkEstimationDispatched< + GROUP_SIZE, HEAD_DIM_QK, POS_ENCODING_MODE, AttentionVariant, Params>; + cudaError_t status = DecodePlan( + static_cast(float_workspace_buffer.data_ptr()), float_workspace_size_in_bytes, + static_cast(int_workspace_buffer.data_ptr()), + static_cast(page_locked_int_workspace_buffer.data_ptr()), + int_workspace_size_in_bytes, plan_info, static_cast(indptr.data_ptr()), + batch_size, num_qo_heads, page_size, enable_cuda_graph, + /*stream=*/stream, work_estimation_func); + + TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); + }); + + return vec_to_tensor(plan_info.ToVector()); +} + +void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, + at::Tensor paged_v_cache, at::Tensor paged_kv_indptr, + at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, + int64_t kv_layout_code, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + DecodePlanInfo plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + QKVLayout kv_layout = static_cast(kv_layout_code); + auto device = q.device(); + int64_t batch_size = q.size(0); + int64_t num_qo_heads = q.size(1); + int64_t num_kv_heads, page_size; + + if (kv_layout == QKVLayout::kHND) { + num_kv_heads = paged_k_cache.size(1); + page_size = paged_k_cache.size(2); + } else { + page_size = paged_k_cache.size(1); + num_kv_heads = paged_k_cache.size(2); + } + uint32_t head_dim_qk = q.size(2); + uint32_t head_dim_vo = paged_v_cache.size(3); + + TORCH_CHECK(head_dim_qk == head_dim_vo, + "CUDA cores template only supports equal head dim for QK and VO, please use tensor " + "cores template for different head dim"); + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == batch_size, lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == num_qo_heads, lse.size(1), q.size(1)); + } + + void* float_buffer = static_cast(float_workspace_buffer.data_ptr()); + void* int_buffer = static_cast(int_workspace_buffer.data_ptr()); + + // get q_scalar_type and kv_scalar_type + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = paged_k_cache.scalar_type(); + + // get q_stride_n and q_stride_h + const auto q_stride_n = q.stride(0); + const auto q_stride_h = q.stride(1); + + // get kv_cache_strides + const int64_t* kv_cache_strides = nullptr; + auto k_strides = paged_k_cache.strides(); + auto v_strides = paged_v_cache.strides(); + TORCH_CHECK(k_strides == v_strides, "k/v strides must be identical"); + kv_cache_strides = k_strides.data(); + + const c10::cuda::OptionalCUDAGuard device_guard(device); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + paged_kv_t paged_kv( + num_kv_heads, page_size, HEAD_DIM_QK, batch_size, kv_layout, + static_cast(paged_k_cache.data_ptr()), + static_cast(paged_v_cache.data_ptr()), kv_cache_strides, + static_cast(paged_kv_indices.data_ptr()), + static_cast(paged_kv_indptr.data_ptr()), + static_cast(paged_kv_last_page_len.data_ptr())); + + Params params; + params.q = static_cast(q.data_ptr()); + params.paged_kv = paged_kv; + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.padded_batch_size = 0; + params.num_qo_heads = num_qo_heads; + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.window_left = window_left; + params.request_indices = nullptr; + params.kv_tile_indices = nullptr; + params.o_indptr = nullptr; + params.kv_chunk_size_ptr = nullptr; + params.block_valid_mask = nullptr; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + DTypeO* tmp_v = nullptr; + float* tmp_s = nullptr; + params.request_indices = + GetPtrFromBaseOffset(int_buffer, plan_info.request_indices_offset); + params.kv_tile_indices = + GetPtrFromBaseOffset(int_buffer, plan_info.kv_tile_indices_offset); + params.o_indptr = GetPtrFromBaseOffset(int_buffer, plan_info.o_indptr_offset); + params.kv_chunk_size_ptr = + GetPtrFromBaseOffset(int_buffer, plan_info.kv_chunk_size_ptr_offset); + if (plan_info.split_kv) { + tmp_v = GetPtrFromBaseOffset(float_buffer, plan_info.v_offset); + tmp_s = GetPtrFromBaseOffset(float_buffer, plan_info.s_offset); + if (plan_info.enable_cuda_graph) { + params.block_valid_mask = + GetPtrFromBaseOffset(int_buffer, plan_info.block_valid_mask_offset); + } + } + params.padded_batch_size = plan_info.padded_batch_size; + + cudaError_t status = + flashinfer::BatchDecodeWithPagedKVCacheDispatched(params, tmp_v, + tmp_s, enable_pdl, + /*stream=*/stream); + TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); +} diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode_config.inc b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..aaad98a9c8afc735ce65220a4c24350af7f22bcd --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode_config.inc @@ -0,0 +1,71 @@ +#pragma once +#include +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta +#define ADDITIONAL_PARAMS_SETTER params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, ...) { \ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ +} + +using namespace flashinfer; + +using DTypeQ = nv_bfloat16; +using DTypeKV = __nv_fp8_e4m3; +using DTypeO = nv_bfloat16; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 64; +constexpr int HEAD_DIM_VO = 64; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = true; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + DTypeQ* q; + paged_kv_t paged_kv; + DTypeO* o; + float* lse; + + float* maybe_alibi_slopes; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; + + + uint32_t padded_batch_size; + uint32_t num_qo_heads; + IdType q_stride_n; + IdType q_stride_h; + int32_t window_left; + bool enable_pdl; + + IdType* request_indices; + IdType* kv_tile_indices; + IdType* o_indptr; + IdType* kv_chunk_size_ptr; + bool* block_valid_mask; + bool partition_kv; + + __host__ __device__ __forceinline__ int32_t get_qo_len(int32_t batch_idx) const { return 1; } + + __host__ __device__ __forceinline__ int32_t get_kv_len(int32_t batch_idx) const { + return paged_kv.get_length(batch_idx); + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode_jit_pybind.cu b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..91e944770b8cee57481f0631545cef02694f50cc --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode_jit_pybind.cu @@ -0,0 +1,40 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "batch_decode_config.inc" +#include "pytorch_extension_utils.h" + +at::Tensor BatchDecodeWithPagedKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size, + int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, + int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo, + at::Tensor empty_q_data, at::Tensor empty_kv_data); + +void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, + at::Tensor paged_v_cache, at::Tensor paged_kv_indptr, + at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, + int64_t kv_layout_code, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Batched decode with paged KV-Cache plan + m.def("plan", BatchDecodeWithPagedKVCachePlan); + // Batched decode with paged KV-Cache run + m.def("run", BatchDecodeWithPagedKVCacheRun); +} diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode_kernel.cu b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode_kernel.cu new file mode 100644 index 0000000000000000000000000000000000000000..6cb7d97320a918556f4e46afbaba5444a0334225 --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode_kernel.cu @@ -0,0 +1,13 @@ +#include +#include "batch_decode_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t +BatchDecodeWithPagedKVCacheDispatched<64, PosEncodingMode::kNone, DefaultAttention, Params>( + Params params, nv_bfloat16* tmp_v, + float* tmp_s, bool enable_pdl, cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode.cu b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode.cu new file mode 100644 index 0000000000000000000000000000000000000000..02c75e02718f8f8ce9b6036f8f788d9a9e254cf1 --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode.cu @@ -0,0 +1,197 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include +#include + +#include "batch_decode_config.inc" +#include "pytorch_conversion_utils.h" +#include "pytorch_extension_utils.h" + +namespace flashinfer { + +template +cudaError_t BatchDecodeWithPagedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v, + float* tmp_s, bool enable_pdl, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +at::Tensor BatchDecodeWithPagedKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size, + int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, + int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo, + at::Tensor empty_q_data, at::Tensor empty_kv_data) { + size_t float_workspace_size_in_bytes = + float_workspace_buffer.size(0) * float_workspace_buffer.element_size(); + size_t int_workspace_size_in_bytes = + int_workspace_buffer.size(0) * int_workspace_buffer.element_size(); + + DecodePlanInfo plan_info; + + auto q_scalar_type = empty_q_data.scalar_type(); + auto kv_scalar_type = empty_kv_data.scalar_type(); + + TORCH_CHECK(head_dim_qk == head_dim_vo, + "CUDA cores template only supports equal head dim for QK and VO, please use tensor " + "cores template for different head dim"); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + DISPATCH_GQA_GROUP_SIZE(num_qo_heads / num_kv_heads, GROUP_SIZE, { + auto work_estimation_func = BatchDecodeWithPagedKVCacheWorkEstimationDispatched< + GROUP_SIZE, HEAD_DIM_QK, POS_ENCODING_MODE, AttentionVariant, Params>; + cudaError_t status = DecodePlan( + static_cast(float_workspace_buffer.data_ptr()), float_workspace_size_in_bytes, + static_cast(int_workspace_buffer.data_ptr()), + static_cast(page_locked_int_workspace_buffer.data_ptr()), + int_workspace_size_in_bytes, plan_info, static_cast(indptr.data_ptr()), + batch_size, num_qo_heads, page_size, enable_cuda_graph, + /*stream=*/stream, work_estimation_func); + + TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); + }); + + return vec_to_tensor(plan_info.ToVector()); +} + +void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, + at::Tensor paged_v_cache, at::Tensor paged_kv_indptr, + at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, + int64_t kv_layout_code, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + DecodePlanInfo plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + QKVLayout kv_layout = static_cast(kv_layout_code); + auto device = q.device(); + int64_t batch_size = q.size(0); + int64_t num_qo_heads = q.size(1); + int64_t num_kv_heads, page_size; + + if (kv_layout == QKVLayout::kHND) { + num_kv_heads = paged_k_cache.size(1); + page_size = paged_k_cache.size(2); + } else { + page_size = paged_k_cache.size(1); + num_kv_heads = paged_k_cache.size(2); + } + uint32_t head_dim_qk = q.size(2); + uint32_t head_dim_vo = paged_v_cache.size(3); + + TORCH_CHECK(head_dim_qk == head_dim_vo, + "CUDA cores template only supports equal head dim for QK and VO, please use tensor " + "cores template for different head dim"); + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == batch_size, lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == num_qo_heads, lse.size(1), q.size(1)); + } + + void* float_buffer = static_cast(float_workspace_buffer.data_ptr()); + void* int_buffer = static_cast(int_workspace_buffer.data_ptr()); + + // get q_scalar_type and kv_scalar_type + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = paged_k_cache.scalar_type(); + + // get q_stride_n and q_stride_h + const auto q_stride_n = q.stride(0); + const auto q_stride_h = q.stride(1); + + // get kv_cache_strides + const int64_t* kv_cache_strides = nullptr; + auto k_strides = paged_k_cache.strides(); + auto v_strides = paged_v_cache.strides(); + TORCH_CHECK(k_strides == v_strides, "k/v strides must be identical"); + kv_cache_strides = k_strides.data(); + + const c10::cuda::OptionalCUDAGuard device_guard(device); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + paged_kv_t paged_kv( + num_kv_heads, page_size, HEAD_DIM_QK, batch_size, kv_layout, + static_cast(paged_k_cache.data_ptr()), + static_cast(paged_v_cache.data_ptr()), kv_cache_strides, + static_cast(paged_kv_indices.data_ptr()), + static_cast(paged_kv_indptr.data_ptr()), + static_cast(paged_kv_last_page_len.data_ptr())); + + Params params; + params.q = static_cast(q.data_ptr()); + params.paged_kv = paged_kv; + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.padded_batch_size = 0; + params.num_qo_heads = num_qo_heads; + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.window_left = window_left; + params.request_indices = nullptr; + params.kv_tile_indices = nullptr; + params.o_indptr = nullptr; + params.kv_chunk_size_ptr = nullptr; + params.block_valid_mask = nullptr; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + DTypeO* tmp_v = nullptr; + float* tmp_s = nullptr; + params.request_indices = + GetPtrFromBaseOffset(int_buffer, plan_info.request_indices_offset); + params.kv_tile_indices = + GetPtrFromBaseOffset(int_buffer, plan_info.kv_tile_indices_offset); + params.o_indptr = GetPtrFromBaseOffset(int_buffer, plan_info.o_indptr_offset); + params.kv_chunk_size_ptr = + GetPtrFromBaseOffset(int_buffer, plan_info.kv_chunk_size_ptr_offset); + if (plan_info.split_kv) { + tmp_v = GetPtrFromBaseOffset(float_buffer, plan_info.v_offset); + tmp_s = GetPtrFromBaseOffset(float_buffer, plan_info.s_offset); + if (plan_info.enable_cuda_graph) { + params.block_valid_mask = + GetPtrFromBaseOffset(int_buffer, plan_info.block_valid_mask_offset); + } + } + params.padded_batch_size = plan_info.padded_batch_size; + + cudaError_t status = + flashinfer::BatchDecodeWithPagedKVCacheDispatched(params, tmp_v, + tmp_s, enable_pdl, + /*stream=*/stream); + TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); +} diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_config.inc b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..8a9327e0f6bca058ce1fc1c5552a401d58384218 --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_config.inc @@ -0,0 +1,71 @@ +#pragma once +#include +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta +#define ADDITIONAL_PARAMS_SETTER params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, ...) { \ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ +} + +using namespace flashinfer; + +using DTypeQ = half; +using DTypeKV = __nv_fp8_e4m3; +using DTypeO = half; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 128; +constexpr int HEAD_DIM_VO = 128; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = false; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + DTypeQ* q; + paged_kv_t paged_kv; + DTypeO* o; + float* lse; + + float* maybe_alibi_slopes; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; + + + uint32_t padded_batch_size; + uint32_t num_qo_heads; + IdType q_stride_n; + IdType q_stride_h; + int32_t window_left; + bool enable_pdl; + + IdType* request_indices; + IdType* kv_tile_indices; + IdType* o_indptr; + IdType* kv_chunk_size_ptr; + bool* block_valid_mask; + bool partition_kv; + + __host__ __device__ __forceinline__ int32_t get_qo_len(int32_t batch_idx) const { return 1; } + + __host__ __device__ __forceinline__ int32_t get_kv_len(int32_t batch_idx) const { + return paged_kv.get_length(batch_idx); + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_jit_pybind.cu b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..91e944770b8cee57481f0631545cef02694f50cc --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_jit_pybind.cu @@ -0,0 +1,40 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "batch_decode_config.inc" +#include "pytorch_extension_utils.h" + +at::Tensor BatchDecodeWithPagedKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size, + int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, + int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo, + at::Tensor empty_q_data, at::Tensor empty_kv_data); + +void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, + at::Tensor paged_v_cache, at::Tensor paged_kv_indptr, + at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, + int64_t kv_layout_code, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Batched decode with paged KV-Cache plan + m.def("plan", BatchDecodeWithPagedKVCachePlan); + // Batched decode with paged KV-Cache run + m.def("run", BatchDecodeWithPagedKVCacheRun); +} diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_kernel.cu b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_kernel.cu new file mode 100644 index 0000000000000000000000000000000000000000..2c6c3ba12784bab5e678073e6aa417f9df22d684 --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_kernel.cu @@ -0,0 +1,13 @@ +#include +#include "batch_decode_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t +BatchDecodeWithPagedKVCacheDispatched<128, PosEncodingMode::kNone, DefaultAttention, Params>( + Params params, half* tmp_v, + float* tmp_s, bool enable_pdl, cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode.cu b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode.cu new file mode 100644 index 0000000000000000000000000000000000000000..02c75e02718f8f8ce9b6036f8f788d9a9e254cf1 --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode.cu @@ -0,0 +1,197 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include +#include + +#include "batch_decode_config.inc" +#include "pytorch_conversion_utils.h" +#include "pytorch_extension_utils.h" + +namespace flashinfer { + +template +cudaError_t BatchDecodeWithPagedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v, + float* tmp_s, bool enable_pdl, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +at::Tensor BatchDecodeWithPagedKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size, + int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, + int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo, + at::Tensor empty_q_data, at::Tensor empty_kv_data) { + size_t float_workspace_size_in_bytes = + float_workspace_buffer.size(0) * float_workspace_buffer.element_size(); + size_t int_workspace_size_in_bytes = + int_workspace_buffer.size(0) * int_workspace_buffer.element_size(); + + DecodePlanInfo plan_info; + + auto q_scalar_type = empty_q_data.scalar_type(); + auto kv_scalar_type = empty_kv_data.scalar_type(); + + TORCH_CHECK(head_dim_qk == head_dim_vo, + "CUDA cores template only supports equal head dim for QK and VO, please use tensor " + "cores template for different head dim"); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + DISPATCH_GQA_GROUP_SIZE(num_qo_heads / num_kv_heads, GROUP_SIZE, { + auto work_estimation_func = BatchDecodeWithPagedKVCacheWorkEstimationDispatched< + GROUP_SIZE, HEAD_DIM_QK, POS_ENCODING_MODE, AttentionVariant, Params>; + cudaError_t status = DecodePlan( + static_cast(float_workspace_buffer.data_ptr()), float_workspace_size_in_bytes, + static_cast(int_workspace_buffer.data_ptr()), + static_cast(page_locked_int_workspace_buffer.data_ptr()), + int_workspace_size_in_bytes, plan_info, static_cast(indptr.data_ptr()), + batch_size, num_qo_heads, page_size, enable_cuda_graph, + /*stream=*/stream, work_estimation_func); + + TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); + }); + + return vec_to_tensor(plan_info.ToVector()); +} + +void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, + at::Tensor paged_v_cache, at::Tensor paged_kv_indptr, + at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, + int64_t kv_layout_code, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + DecodePlanInfo plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + QKVLayout kv_layout = static_cast(kv_layout_code); + auto device = q.device(); + int64_t batch_size = q.size(0); + int64_t num_qo_heads = q.size(1); + int64_t num_kv_heads, page_size; + + if (kv_layout == QKVLayout::kHND) { + num_kv_heads = paged_k_cache.size(1); + page_size = paged_k_cache.size(2); + } else { + page_size = paged_k_cache.size(1); + num_kv_heads = paged_k_cache.size(2); + } + uint32_t head_dim_qk = q.size(2); + uint32_t head_dim_vo = paged_v_cache.size(3); + + TORCH_CHECK(head_dim_qk == head_dim_vo, + "CUDA cores template only supports equal head dim for QK and VO, please use tensor " + "cores template for different head dim"); + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == batch_size, lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == num_qo_heads, lse.size(1), q.size(1)); + } + + void* float_buffer = static_cast(float_workspace_buffer.data_ptr()); + void* int_buffer = static_cast(int_workspace_buffer.data_ptr()); + + // get q_scalar_type and kv_scalar_type + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = paged_k_cache.scalar_type(); + + // get q_stride_n and q_stride_h + const auto q_stride_n = q.stride(0); + const auto q_stride_h = q.stride(1); + + // get kv_cache_strides + const int64_t* kv_cache_strides = nullptr; + auto k_strides = paged_k_cache.strides(); + auto v_strides = paged_v_cache.strides(); + TORCH_CHECK(k_strides == v_strides, "k/v strides must be identical"); + kv_cache_strides = k_strides.data(); + + const c10::cuda::OptionalCUDAGuard device_guard(device); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + paged_kv_t paged_kv( + num_kv_heads, page_size, HEAD_DIM_QK, batch_size, kv_layout, + static_cast(paged_k_cache.data_ptr()), + static_cast(paged_v_cache.data_ptr()), kv_cache_strides, + static_cast(paged_kv_indices.data_ptr()), + static_cast(paged_kv_indptr.data_ptr()), + static_cast(paged_kv_last_page_len.data_ptr())); + + Params params; + params.q = static_cast(q.data_ptr()); + params.paged_kv = paged_kv; + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.padded_batch_size = 0; + params.num_qo_heads = num_qo_heads; + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.window_left = window_left; + params.request_indices = nullptr; + params.kv_tile_indices = nullptr; + params.o_indptr = nullptr; + params.kv_chunk_size_ptr = nullptr; + params.block_valid_mask = nullptr; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + DTypeO* tmp_v = nullptr; + float* tmp_s = nullptr; + params.request_indices = + GetPtrFromBaseOffset(int_buffer, plan_info.request_indices_offset); + params.kv_tile_indices = + GetPtrFromBaseOffset(int_buffer, plan_info.kv_tile_indices_offset); + params.o_indptr = GetPtrFromBaseOffset(int_buffer, plan_info.o_indptr_offset); + params.kv_chunk_size_ptr = + GetPtrFromBaseOffset(int_buffer, plan_info.kv_chunk_size_ptr_offset); + if (plan_info.split_kv) { + tmp_v = GetPtrFromBaseOffset(float_buffer, plan_info.v_offset); + tmp_s = GetPtrFromBaseOffset(float_buffer, plan_info.s_offset); + if (plan_info.enable_cuda_graph) { + params.block_valid_mask = + GetPtrFromBaseOffset(int_buffer, plan_info.block_valid_mask_offset); + } + } + params.padded_batch_size = plan_info.padded_batch_size; + + cudaError_t status = + flashinfer::BatchDecodeWithPagedKVCacheDispatched(params, tmp_v, + tmp_s, enable_pdl, + /*stream=*/stream); + TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); +} diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode_config.inc b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..98f25964351e920a56c287e3f05b5123d71fb67d --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode_config.inc @@ -0,0 +1,71 @@ +#pragma once +#include +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta +#define ADDITIONAL_PARAMS_SETTER params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, ...) { \ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ +} + +using namespace flashinfer; + +using DTypeQ = half; +using DTypeKV = __nv_fp8_e4m3; +using DTypeO = half; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 256; +constexpr int HEAD_DIM_VO = 256; +constexpr auto USE_LOGITS_SOFT_CAP = true; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = true; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + DTypeQ* q; + paged_kv_t paged_kv; + DTypeO* o; + float* lse; + + float* maybe_alibi_slopes; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; + + + uint32_t padded_batch_size; + uint32_t num_qo_heads; + IdType q_stride_n; + IdType q_stride_h; + int32_t window_left; + bool enable_pdl; + + IdType* request_indices; + IdType* kv_tile_indices; + IdType* o_indptr; + IdType* kv_chunk_size_ptr; + bool* block_valid_mask; + bool partition_kv; + + __host__ __device__ __forceinline__ int32_t get_qo_len(int32_t batch_idx) const { return 1; } + + __host__ __device__ __forceinline__ int32_t get_kv_len(int32_t batch_idx) const { + return paged_kv.get_length(batch_idx); + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode_jit_pybind.cu b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..91e944770b8cee57481f0631545cef02694f50cc --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode_jit_pybind.cu @@ -0,0 +1,40 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "batch_decode_config.inc" +#include "pytorch_extension_utils.h" + +at::Tensor BatchDecodeWithPagedKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size, + int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, + int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo, + at::Tensor empty_q_data, at::Tensor empty_kv_data); + +void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, + at::Tensor paged_v_cache, at::Tensor paged_kv_indptr, + at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, + int64_t kv_layout_code, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Batched decode with paged KV-Cache plan + m.def("plan", BatchDecodeWithPagedKVCachePlan); + // Batched decode with paged KV-Cache run + m.def("run", BatchDecodeWithPagedKVCacheRun); +} diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode_kernel.cu b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode_kernel.cu new file mode 100644 index 0000000000000000000000000000000000000000..fc08143c5b0cc6d1ba272d41bae81db75865298c --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode_kernel.cu @@ -0,0 +1,13 @@ +#include +#include "batch_decode_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t +BatchDecodeWithPagedKVCacheDispatched<256, PosEncodingMode::kNone, DefaultAttention, Params>( + Params params, half* tmp_v, + float* tmp_s, bool enable_pdl, cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode.cu b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode.cu new file mode 100644 index 0000000000000000000000000000000000000000..02c75e02718f8f8ce9b6036f8f788d9a9e254cf1 --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode.cu @@ -0,0 +1,197 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include +#include + +#include "batch_decode_config.inc" +#include "pytorch_conversion_utils.h" +#include "pytorch_extension_utils.h" + +namespace flashinfer { + +template +cudaError_t BatchDecodeWithPagedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v, + float* tmp_s, bool enable_pdl, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +at::Tensor BatchDecodeWithPagedKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size, + int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, + int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo, + at::Tensor empty_q_data, at::Tensor empty_kv_data) { + size_t float_workspace_size_in_bytes = + float_workspace_buffer.size(0) * float_workspace_buffer.element_size(); + size_t int_workspace_size_in_bytes = + int_workspace_buffer.size(0) * int_workspace_buffer.element_size(); + + DecodePlanInfo plan_info; + + auto q_scalar_type = empty_q_data.scalar_type(); + auto kv_scalar_type = empty_kv_data.scalar_type(); + + TORCH_CHECK(head_dim_qk == head_dim_vo, + "CUDA cores template only supports equal head dim for QK and VO, please use tensor " + "cores template for different head dim"); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + DISPATCH_GQA_GROUP_SIZE(num_qo_heads / num_kv_heads, GROUP_SIZE, { + auto work_estimation_func = BatchDecodeWithPagedKVCacheWorkEstimationDispatched< + GROUP_SIZE, HEAD_DIM_QK, POS_ENCODING_MODE, AttentionVariant, Params>; + cudaError_t status = DecodePlan( + static_cast(float_workspace_buffer.data_ptr()), float_workspace_size_in_bytes, + static_cast(int_workspace_buffer.data_ptr()), + static_cast(page_locked_int_workspace_buffer.data_ptr()), + int_workspace_size_in_bytes, plan_info, static_cast(indptr.data_ptr()), + batch_size, num_qo_heads, page_size, enable_cuda_graph, + /*stream=*/stream, work_estimation_func); + + TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); + }); + + return vec_to_tensor(plan_info.ToVector()); +} + +void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, + at::Tensor paged_v_cache, at::Tensor paged_kv_indptr, + at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, + int64_t kv_layout_code, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + DecodePlanInfo plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + QKVLayout kv_layout = static_cast(kv_layout_code); + auto device = q.device(); + int64_t batch_size = q.size(0); + int64_t num_qo_heads = q.size(1); + int64_t num_kv_heads, page_size; + + if (kv_layout == QKVLayout::kHND) { + num_kv_heads = paged_k_cache.size(1); + page_size = paged_k_cache.size(2); + } else { + page_size = paged_k_cache.size(1); + num_kv_heads = paged_k_cache.size(2); + } + uint32_t head_dim_qk = q.size(2); + uint32_t head_dim_vo = paged_v_cache.size(3); + + TORCH_CHECK(head_dim_qk == head_dim_vo, + "CUDA cores template only supports equal head dim for QK and VO, please use tensor " + "cores template for different head dim"); + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == batch_size, lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == num_qo_heads, lse.size(1), q.size(1)); + } + + void* float_buffer = static_cast(float_workspace_buffer.data_ptr()); + void* int_buffer = static_cast(int_workspace_buffer.data_ptr()); + + // get q_scalar_type and kv_scalar_type + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = paged_k_cache.scalar_type(); + + // get q_stride_n and q_stride_h + const auto q_stride_n = q.stride(0); + const auto q_stride_h = q.stride(1); + + // get kv_cache_strides + const int64_t* kv_cache_strides = nullptr; + auto k_strides = paged_k_cache.strides(); + auto v_strides = paged_v_cache.strides(); + TORCH_CHECK(k_strides == v_strides, "k/v strides must be identical"); + kv_cache_strides = k_strides.data(); + + const c10::cuda::OptionalCUDAGuard device_guard(device); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + paged_kv_t paged_kv( + num_kv_heads, page_size, HEAD_DIM_QK, batch_size, kv_layout, + static_cast(paged_k_cache.data_ptr()), + static_cast(paged_v_cache.data_ptr()), kv_cache_strides, + static_cast(paged_kv_indices.data_ptr()), + static_cast(paged_kv_indptr.data_ptr()), + static_cast(paged_kv_last_page_len.data_ptr())); + + Params params; + params.q = static_cast(q.data_ptr()); + params.paged_kv = paged_kv; + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.padded_batch_size = 0; + params.num_qo_heads = num_qo_heads; + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.window_left = window_left; + params.request_indices = nullptr; + params.kv_tile_indices = nullptr; + params.o_indptr = nullptr; + params.kv_chunk_size_ptr = nullptr; + params.block_valid_mask = nullptr; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + DTypeO* tmp_v = nullptr; + float* tmp_s = nullptr; + params.request_indices = + GetPtrFromBaseOffset(int_buffer, plan_info.request_indices_offset); + params.kv_tile_indices = + GetPtrFromBaseOffset(int_buffer, plan_info.kv_tile_indices_offset); + params.o_indptr = GetPtrFromBaseOffset(int_buffer, plan_info.o_indptr_offset); + params.kv_chunk_size_ptr = + GetPtrFromBaseOffset(int_buffer, plan_info.kv_chunk_size_ptr_offset); + if (plan_info.split_kv) { + tmp_v = GetPtrFromBaseOffset(float_buffer, plan_info.v_offset); + tmp_s = GetPtrFromBaseOffset(float_buffer, plan_info.s_offset); + if (plan_info.enable_cuda_graph) { + params.block_valid_mask = + GetPtrFromBaseOffset(int_buffer, plan_info.block_valid_mask_offset); + } + } + params.padded_batch_size = plan_info.padded_batch_size; + + cudaError_t status = + flashinfer::BatchDecodeWithPagedKVCacheDispatched(params, tmp_v, + tmp_s, enable_pdl, + /*stream=*/stream); + TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); +} diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_config.inc b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..73e2ce5148559dd5900216012ba0ea60aa8ce90d --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_config.inc @@ -0,0 +1,71 @@ +#pragma once +#include +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta +#define ADDITIONAL_PARAMS_SETTER params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, ...) { \ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ +} + +using namespace flashinfer; + +using DTypeQ = half; +using DTypeKV = __nv_fp8_e4m3; +using DTypeO = half; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 64; +constexpr int HEAD_DIM_VO = 64; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = false; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + DTypeQ* q; + paged_kv_t paged_kv; + DTypeO* o; + float* lse; + + float* maybe_alibi_slopes; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; + + + uint32_t padded_batch_size; + uint32_t num_qo_heads; + IdType q_stride_n; + IdType q_stride_h; + int32_t window_left; + bool enable_pdl; + + IdType* request_indices; + IdType* kv_tile_indices; + IdType* o_indptr; + IdType* kv_chunk_size_ptr; + bool* block_valid_mask; + bool partition_kv; + + __host__ __device__ __forceinline__ int32_t get_qo_len(int32_t batch_idx) const { return 1; } + + __host__ __device__ __forceinline__ int32_t get_kv_len(int32_t batch_idx) const { + return paged_kv.get_length(batch_idx); + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_jit_pybind.cu b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..91e944770b8cee57481f0631545cef02694f50cc --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_jit_pybind.cu @@ -0,0 +1,40 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "batch_decode_config.inc" +#include "pytorch_extension_utils.h" + +at::Tensor BatchDecodeWithPagedKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size, + int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, + int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo, + at::Tensor empty_q_data, at::Tensor empty_kv_data); + +void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, + at::Tensor paged_v_cache, at::Tensor paged_kv_indptr, + at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, + int64_t kv_layout_code, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Batched decode with paged KV-Cache plan + m.def("plan", BatchDecodeWithPagedKVCachePlan); + // Batched decode with paged KV-Cache run + m.def("run", BatchDecodeWithPagedKVCacheRun); +} diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_kernel.cu b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_kernel.cu new file mode 100644 index 0000000000000000000000000000000000000000..91336af2f35e622c26c8ef0518429688a1839b6a --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_kernel.cu @@ -0,0 +1,13 @@ +#include +#include "batch_decode_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t +BatchDecodeWithPagedKVCacheDispatched<64, PosEncodingMode::kNone, DefaultAttention, Params>( + Params params, half* tmp_v, + float* tmp_s, bool enable_pdl, cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode.cu b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode.cu new file mode 100644 index 0000000000000000000000000000000000000000..02c75e02718f8f8ce9b6036f8f788d9a9e254cf1 --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode.cu @@ -0,0 +1,197 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include +#include + +#include "batch_decode_config.inc" +#include "pytorch_conversion_utils.h" +#include "pytorch_extension_utils.h" + +namespace flashinfer { + +template +cudaError_t BatchDecodeWithPagedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v, + float* tmp_s, bool enable_pdl, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +at::Tensor BatchDecodeWithPagedKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size, + int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, + int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo, + at::Tensor empty_q_data, at::Tensor empty_kv_data) { + size_t float_workspace_size_in_bytes = + float_workspace_buffer.size(0) * float_workspace_buffer.element_size(); + size_t int_workspace_size_in_bytes = + int_workspace_buffer.size(0) * int_workspace_buffer.element_size(); + + DecodePlanInfo plan_info; + + auto q_scalar_type = empty_q_data.scalar_type(); + auto kv_scalar_type = empty_kv_data.scalar_type(); + + TORCH_CHECK(head_dim_qk == head_dim_vo, + "CUDA cores template only supports equal head dim for QK and VO, please use tensor " + "cores template for different head dim"); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + DISPATCH_GQA_GROUP_SIZE(num_qo_heads / num_kv_heads, GROUP_SIZE, { + auto work_estimation_func = BatchDecodeWithPagedKVCacheWorkEstimationDispatched< + GROUP_SIZE, HEAD_DIM_QK, POS_ENCODING_MODE, AttentionVariant, Params>; + cudaError_t status = DecodePlan( + static_cast(float_workspace_buffer.data_ptr()), float_workspace_size_in_bytes, + static_cast(int_workspace_buffer.data_ptr()), + static_cast(page_locked_int_workspace_buffer.data_ptr()), + int_workspace_size_in_bytes, plan_info, static_cast(indptr.data_ptr()), + batch_size, num_qo_heads, page_size, enable_cuda_graph, + /*stream=*/stream, work_estimation_func); + + TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); + }); + + return vec_to_tensor(plan_info.ToVector()); +} + +void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, + at::Tensor paged_v_cache, at::Tensor paged_kv_indptr, + at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, + int64_t kv_layout_code, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + DecodePlanInfo plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + QKVLayout kv_layout = static_cast(kv_layout_code); + auto device = q.device(); + int64_t batch_size = q.size(0); + int64_t num_qo_heads = q.size(1); + int64_t num_kv_heads, page_size; + + if (kv_layout == QKVLayout::kHND) { + num_kv_heads = paged_k_cache.size(1); + page_size = paged_k_cache.size(2); + } else { + page_size = paged_k_cache.size(1); + num_kv_heads = paged_k_cache.size(2); + } + uint32_t head_dim_qk = q.size(2); + uint32_t head_dim_vo = paged_v_cache.size(3); + + TORCH_CHECK(head_dim_qk == head_dim_vo, + "CUDA cores template only supports equal head dim for QK and VO, please use tensor " + "cores template for different head dim"); + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == batch_size, lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == num_qo_heads, lse.size(1), q.size(1)); + } + + void* float_buffer = static_cast(float_workspace_buffer.data_ptr()); + void* int_buffer = static_cast(int_workspace_buffer.data_ptr()); + + // get q_scalar_type and kv_scalar_type + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = paged_k_cache.scalar_type(); + + // get q_stride_n and q_stride_h + const auto q_stride_n = q.stride(0); + const auto q_stride_h = q.stride(1); + + // get kv_cache_strides + const int64_t* kv_cache_strides = nullptr; + auto k_strides = paged_k_cache.strides(); + auto v_strides = paged_v_cache.strides(); + TORCH_CHECK(k_strides == v_strides, "k/v strides must be identical"); + kv_cache_strides = k_strides.data(); + + const c10::cuda::OptionalCUDAGuard device_guard(device); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + paged_kv_t paged_kv( + num_kv_heads, page_size, HEAD_DIM_QK, batch_size, kv_layout, + static_cast(paged_k_cache.data_ptr()), + static_cast(paged_v_cache.data_ptr()), kv_cache_strides, + static_cast(paged_kv_indices.data_ptr()), + static_cast(paged_kv_indptr.data_ptr()), + static_cast(paged_kv_last_page_len.data_ptr())); + + Params params; + params.q = static_cast(q.data_ptr()); + params.paged_kv = paged_kv; + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.padded_batch_size = 0; + params.num_qo_heads = num_qo_heads; + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.window_left = window_left; + params.request_indices = nullptr; + params.kv_tile_indices = nullptr; + params.o_indptr = nullptr; + params.kv_chunk_size_ptr = nullptr; + params.block_valid_mask = nullptr; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + DTypeO* tmp_v = nullptr; + float* tmp_s = nullptr; + params.request_indices = + GetPtrFromBaseOffset(int_buffer, plan_info.request_indices_offset); + params.kv_tile_indices = + GetPtrFromBaseOffset(int_buffer, plan_info.kv_tile_indices_offset); + params.o_indptr = GetPtrFromBaseOffset(int_buffer, plan_info.o_indptr_offset); + params.kv_chunk_size_ptr = + GetPtrFromBaseOffset(int_buffer, plan_info.kv_chunk_size_ptr_offset); + if (plan_info.split_kv) { + tmp_v = GetPtrFromBaseOffset(float_buffer, plan_info.v_offset); + tmp_s = GetPtrFromBaseOffset(float_buffer, plan_info.s_offset); + if (plan_info.enable_cuda_graph) { + params.block_valid_mask = + GetPtrFromBaseOffset(int_buffer, plan_info.block_valid_mask_offset); + } + } + params.padded_batch_size = plan_info.padded_batch_size; + + cudaError_t status = + flashinfer::BatchDecodeWithPagedKVCacheDispatched(params, tmp_v, + tmp_s, enable_pdl, + /*stream=*/stream); + TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); +} diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode_config.inc b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..5f49d962b27bbb911dd370a94cd85a3bf596d4b0 --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode_config.inc @@ -0,0 +1,71 @@ +#pragma once +#include +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta +#define ADDITIONAL_PARAMS_SETTER params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, ...) { \ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ +} + +using namespace flashinfer; + +using DTypeQ = half; +using DTypeKV = __nv_fp8_e4m3; +using DTypeO = half; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 64; +constexpr int HEAD_DIM_VO = 64; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = true; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + DTypeQ* q; + paged_kv_t paged_kv; + DTypeO* o; + float* lse; + + float* maybe_alibi_slopes; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; + + + uint32_t padded_batch_size; + uint32_t num_qo_heads; + IdType q_stride_n; + IdType q_stride_h; + int32_t window_left; + bool enable_pdl; + + IdType* request_indices; + IdType* kv_tile_indices; + IdType* o_indptr; + IdType* kv_chunk_size_ptr; + bool* block_valid_mask; + bool partition_kv; + + __host__ __device__ __forceinline__ int32_t get_qo_len(int32_t batch_idx) const { return 1; } + + __host__ __device__ __forceinline__ int32_t get_kv_len(int32_t batch_idx) const { + return paged_kv.get_length(batch_idx); + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode_jit_pybind.cu b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..91e944770b8cee57481f0631545cef02694f50cc --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode_jit_pybind.cu @@ -0,0 +1,40 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "batch_decode_config.inc" +#include "pytorch_extension_utils.h" + +at::Tensor BatchDecodeWithPagedKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size, + int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, + int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo, + at::Tensor empty_q_data, at::Tensor empty_kv_data); + +void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, + at::Tensor paged_v_cache, at::Tensor paged_kv_indptr, + at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, + int64_t kv_layout_code, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Batched decode with paged KV-Cache plan + m.def("plan", BatchDecodeWithPagedKVCachePlan); + // Batched decode with paged KV-Cache run + m.def("run", BatchDecodeWithPagedKVCacheRun); +} diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode_kernel.cu b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode_kernel.cu new file mode 100644 index 0000000000000000000000000000000000000000..f9cecd66f3d2b259979d5fb2b612e9212bb3b57a --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode_kernel.cu @@ -0,0 +1,13 @@ +#include +#include "batch_decode_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t +BatchDecodeWithPagedKVCacheDispatched<64, PosEncodingMode::kNone, DefaultAttention, Params>( + Params params, half* tmp_v, + float* tmp_s, bool enable_pdl, cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode.cu b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode.cu new file mode 100644 index 0000000000000000000000000000000000000000..02c75e02718f8f8ce9b6036f8f788d9a9e254cf1 --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode.cu @@ -0,0 +1,197 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include +#include + +#include "batch_decode_config.inc" +#include "pytorch_conversion_utils.h" +#include "pytorch_extension_utils.h" + +namespace flashinfer { + +template +cudaError_t BatchDecodeWithPagedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v, + float* tmp_s, bool enable_pdl, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +at::Tensor BatchDecodeWithPagedKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size, + int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, + int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo, + at::Tensor empty_q_data, at::Tensor empty_kv_data) { + size_t float_workspace_size_in_bytes = + float_workspace_buffer.size(0) * float_workspace_buffer.element_size(); + size_t int_workspace_size_in_bytes = + int_workspace_buffer.size(0) * int_workspace_buffer.element_size(); + + DecodePlanInfo plan_info; + + auto q_scalar_type = empty_q_data.scalar_type(); + auto kv_scalar_type = empty_kv_data.scalar_type(); + + TORCH_CHECK(head_dim_qk == head_dim_vo, + "CUDA cores template only supports equal head dim for QK and VO, please use tensor " + "cores template for different head dim"); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + DISPATCH_GQA_GROUP_SIZE(num_qo_heads / num_kv_heads, GROUP_SIZE, { + auto work_estimation_func = BatchDecodeWithPagedKVCacheWorkEstimationDispatched< + GROUP_SIZE, HEAD_DIM_QK, POS_ENCODING_MODE, AttentionVariant, Params>; + cudaError_t status = DecodePlan( + static_cast(float_workspace_buffer.data_ptr()), float_workspace_size_in_bytes, + static_cast(int_workspace_buffer.data_ptr()), + static_cast(page_locked_int_workspace_buffer.data_ptr()), + int_workspace_size_in_bytes, plan_info, static_cast(indptr.data_ptr()), + batch_size, num_qo_heads, page_size, enable_cuda_graph, + /*stream=*/stream, work_estimation_func); + + TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); + }); + + return vec_to_tensor(plan_info.ToVector()); +} + +void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, + at::Tensor paged_v_cache, at::Tensor paged_kv_indptr, + at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, + int64_t kv_layout_code, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + DecodePlanInfo plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + QKVLayout kv_layout = static_cast(kv_layout_code); + auto device = q.device(); + int64_t batch_size = q.size(0); + int64_t num_qo_heads = q.size(1); + int64_t num_kv_heads, page_size; + + if (kv_layout == QKVLayout::kHND) { + num_kv_heads = paged_k_cache.size(1); + page_size = paged_k_cache.size(2); + } else { + page_size = paged_k_cache.size(1); + num_kv_heads = paged_k_cache.size(2); + } + uint32_t head_dim_qk = q.size(2); + uint32_t head_dim_vo = paged_v_cache.size(3); + + TORCH_CHECK(head_dim_qk == head_dim_vo, + "CUDA cores template only supports equal head dim for QK and VO, please use tensor " + "cores template for different head dim"); + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == batch_size, lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == num_qo_heads, lse.size(1), q.size(1)); + } + + void* float_buffer = static_cast(float_workspace_buffer.data_ptr()); + void* int_buffer = static_cast(int_workspace_buffer.data_ptr()); + + // get q_scalar_type and kv_scalar_type + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = paged_k_cache.scalar_type(); + + // get q_stride_n and q_stride_h + const auto q_stride_n = q.stride(0); + const auto q_stride_h = q.stride(1); + + // get kv_cache_strides + const int64_t* kv_cache_strides = nullptr; + auto k_strides = paged_k_cache.strides(); + auto v_strides = paged_v_cache.strides(); + TORCH_CHECK(k_strides == v_strides, "k/v strides must be identical"); + kv_cache_strides = k_strides.data(); + + const c10::cuda::OptionalCUDAGuard device_guard(device); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + paged_kv_t paged_kv( + num_kv_heads, page_size, HEAD_DIM_QK, batch_size, kv_layout, + static_cast(paged_k_cache.data_ptr()), + static_cast(paged_v_cache.data_ptr()), kv_cache_strides, + static_cast(paged_kv_indices.data_ptr()), + static_cast(paged_kv_indptr.data_ptr()), + static_cast(paged_kv_last_page_len.data_ptr())); + + Params params; + params.q = static_cast(q.data_ptr()); + params.paged_kv = paged_kv; + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.padded_batch_size = 0; + params.num_qo_heads = num_qo_heads; + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.window_left = window_left; + params.request_indices = nullptr; + params.kv_tile_indices = nullptr; + params.o_indptr = nullptr; + params.kv_chunk_size_ptr = nullptr; + params.block_valid_mask = nullptr; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + DTypeO* tmp_v = nullptr; + float* tmp_s = nullptr; + params.request_indices = + GetPtrFromBaseOffset(int_buffer, plan_info.request_indices_offset); + params.kv_tile_indices = + GetPtrFromBaseOffset(int_buffer, plan_info.kv_tile_indices_offset); + params.o_indptr = GetPtrFromBaseOffset(int_buffer, plan_info.o_indptr_offset); + params.kv_chunk_size_ptr = + GetPtrFromBaseOffset(int_buffer, plan_info.kv_chunk_size_ptr_offset); + if (plan_info.split_kv) { + tmp_v = GetPtrFromBaseOffset(float_buffer, plan_info.v_offset); + tmp_s = GetPtrFromBaseOffset(float_buffer, plan_info.s_offset); + if (plan_info.enable_cuda_graph) { + params.block_valid_mask = + GetPtrFromBaseOffset(int_buffer, plan_info.block_valid_mask_offset); + } + } + params.padded_batch_size = plan_info.padded_batch_size; + + cudaError_t status = + flashinfer::BatchDecodeWithPagedKVCacheDispatched(params, tmp_v, + tmp_s, enable_pdl, + /*stream=*/stream); + TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); +} diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_config.inc b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..b29fe5afe59b18097e80f8b08e67db13348b4f0a --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_config.inc @@ -0,0 +1,71 @@ +#pragma once +#include +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta +#define ADDITIONAL_PARAMS_SETTER params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, ...) { \ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ +} + +using namespace flashinfer; + +using DTypeQ = half; +using DTypeKV = half; +using DTypeO = half; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 128; +constexpr int HEAD_DIM_VO = 128; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = false; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + DTypeQ* q; + paged_kv_t paged_kv; + DTypeO* o; + float* lse; + + float* maybe_alibi_slopes; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; + + + uint32_t padded_batch_size; + uint32_t num_qo_heads; + IdType q_stride_n; + IdType q_stride_h; + int32_t window_left; + bool enable_pdl; + + IdType* request_indices; + IdType* kv_tile_indices; + IdType* o_indptr; + IdType* kv_chunk_size_ptr; + bool* block_valid_mask; + bool partition_kv; + + __host__ __device__ __forceinline__ int32_t get_qo_len(int32_t batch_idx) const { return 1; } + + __host__ __device__ __forceinline__ int32_t get_kv_len(int32_t batch_idx) const { + return paged_kv.get_length(batch_idx); + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_jit_pybind.cu b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..91e944770b8cee57481f0631545cef02694f50cc --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_jit_pybind.cu @@ -0,0 +1,40 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "batch_decode_config.inc" +#include "pytorch_extension_utils.h" + +at::Tensor BatchDecodeWithPagedKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size, + int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, + int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo, + at::Tensor empty_q_data, at::Tensor empty_kv_data); + +void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, + at::Tensor paged_v_cache, at::Tensor paged_kv_indptr, + at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, + int64_t kv_layout_code, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Batched decode with paged KV-Cache plan + m.def("plan", BatchDecodeWithPagedKVCachePlan); + // Batched decode with paged KV-Cache run + m.def("run", BatchDecodeWithPagedKVCacheRun); +} diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_kernel.cu b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_kernel.cu new file mode 100644 index 0000000000000000000000000000000000000000..2c6c3ba12784bab5e678073e6aa417f9df22d684 --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_kernel.cu @@ -0,0 +1,13 @@ +#include +#include "batch_decode_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t +BatchDecodeWithPagedKVCacheDispatched<128, PosEncodingMode::kNone, DefaultAttention, Params>( + Params params, half* tmp_v, + float* tmp_s, bool enable_pdl, cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode.cu b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode.cu new file mode 100644 index 0000000000000000000000000000000000000000..02c75e02718f8f8ce9b6036f8f788d9a9e254cf1 --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode.cu @@ -0,0 +1,197 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include +#include + +#include "batch_decode_config.inc" +#include "pytorch_conversion_utils.h" +#include "pytorch_extension_utils.h" + +namespace flashinfer { + +template +cudaError_t BatchDecodeWithPagedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v, + float* tmp_s, bool enable_pdl, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +at::Tensor BatchDecodeWithPagedKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size, + int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, + int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo, + at::Tensor empty_q_data, at::Tensor empty_kv_data) { + size_t float_workspace_size_in_bytes = + float_workspace_buffer.size(0) * float_workspace_buffer.element_size(); + size_t int_workspace_size_in_bytes = + int_workspace_buffer.size(0) * int_workspace_buffer.element_size(); + + DecodePlanInfo plan_info; + + auto q_scalar_type = empty_q_data.scalar_type(); + auto kv_scalar_type = empty_kv_data.scalar_type(); + + TORCH_CHECK(head_dim_qk == head_dim_vo, + "CUDA cores template only supports equal head dim for QK and VO, please use tensor " + "cores template for different head dim"); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + DISPATCH_GQA_GROUP_SIZE(num_qo_heads / num_kv_heads, GROUP_SIZE, { + auto work_estimation_func = BatchDecodeWithPagedKVCacheWorkEstimationDispatched< + GROUP_SIZE, HEAD_DIM_QK, POS_ENCODING_MODE, AttentionVariant, Params>; + cudaError_t status = DecodePlan( + static_cast(float_workspace_buffer.data_ptr()), float_workspace_size_in_bytes, + static_cast(int_workspace_buffer.data_ptr()), + static_cast(page_locked_int_workspace_buffer.data_ptr()), + int_workspace_size_in_bytes, plan_info, static_cast(indptr.data_ptr()), + batch_size, num_qo_heads, page_size, enable_cuda_graph, + /*stream=*/stream, work_estimation_func); + + TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); + }); + + return vec_to_tensor(plan_info.ToVector()); +} + +void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, + at::Tensor paged_v_cache, at::Tensor paged_kv_indptr, + at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, + int64_t kv_layout_code, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + DecodePlanInfo plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + QKVLayout kv_layout = static_cast(kv_layout_code); + auto device = q.device(); + int64_t batch_size = q.size(0); + int64_t num_qo_heads = q.size(1); + int64_t num_kv_heads, page_size; + + if (kv_layout == QKVLayout::kHND) { + num_kv_heads = paged_k_cache.size(1); + page_size = paged_k_cache.size(2); + } else { + page_size = paged_k_cache.size(1); + num_kv_heads = paged_k_cache.size(2); + } + uint32_t head_dim_qk = q.size(2); + uint32_t head_dim_vo = paged_v_cache.size(3); + + TORCH_CHECK(head_dim_qk == head_dim_vo, + "CUDA cores template only supports equal head dim for QK and VO, please use tensor " + "cores template for different head dim"); + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == batch_size, lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == num_qo_heads, lse.size(1), q.size(1)); + } + + void* float_buffer = static_cast(float_workspace_buffer.data_ptr()); + void* int_buffer = static_cast(int_workspace_buffer.data_ptr()); + + // get q_scalar_type and kv_scalar_type + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = paged_k_cache.scalar_type(); + + // get q_stride_n and q_stride_h + const auto q_stride_n = q.stride(0); + const auto q_stride_h = q.stride(1); + + // get kv_cache_strides + const int64_t* kv_cache_strides = nullptr; + auto k_strides = paged_k_cache.strides(); + auto v_strides = paged_v_cache.strides(); + TORCH_CHECK(k_strides == v_strides, "k/v strides must be identical"); + kv_cache_strides = k_strides.data(); + + const c10::cuda::OptionalCUDAGuard device_guard(device); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + paged_kv_t paged_kv( + num_kv_heads, page_size, HEAD_DIM_QK, batch_size, kv_layout, + static_cast(paged_k_cache.data_ptr()), + static_cast(paged_v_cache.data_ptr()), kv_cache_strides, + static_cast(paged_kv_indices.data_ptr()), + static_cast(paged_kv_indptr.data_ptr()), + static_cast(paged_kv_last_page_len.data_ptr())); + + Params params; + params.q = static_cast(q.data_ptr()); + params.paged_kv = paged_kv; + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.padded_batch_size = 0; + params.num_qo_heads = num_qo_heads; + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.window_left = window_left; + params.request_indices = nullptr; + params.kv_tile_indices = nullptr; + params.o_indptr = nullptr; + params.kv_chunk_size_ptr = nullptr; + params.block_valid_mask = nullptr; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + DTypeO* tmp_v = nullptr; + float* tmp_s = nullptr; + params.request_indices = + GetPtrFromBaseOffset(int_buffer, plan_info.request_indices_offset); + params.kv_tile_indices = + GetPtrFromBaseOffset(int_buffer, plan_info.kv_tile_indices_offset); + params.o_indptr = GetPtrFromBaseOffset(int_buffer, plan_info.o_indptr_offset); + params.kv_chunk_size_ptr = + GetPtrFromBaseOffset(int_buffer, plan_info.kv_chunk_size_ptr_offset); + if (plan_info.split_kv) { + tmp_v = GetPtrFromBaseOffset(float_buffer, plan_info.v_offset); + tmp_s = GetPtrFromBaseOffset(float_buffer, plan_info.s_offset); + if (plan_info.enable_cuda_graph) { + params.block_valid_mask = + GetPtrFromBaseOffset(int_buffer, plan_info.block_valid_mask_offset); + } + } + params.padded_batch_size = plan_info.padded_batch_size; + + cudaError_t status = + flashinfer::BatchDecodeWithPagedKVCacheDispatched(params, tmp_v, + tmp_s, enable_pdl, + /*stream=*/stream); + TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); +} diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode_config.inc b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..bff612e05a0dc2df177db9f7752fcc8081779b74 --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode_config.inc @@ -0,0 +1,71 @@ +#pragma once +#include +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta +#define ADDITIONAL_PARAMS_SETTER params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, ...) { \ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ +} + +using namespace flashinfer; + +using DTypeQ = half; +using DTypeKV = half; +using DTypeO = half; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 256; +constexpr int HEAD_DIM_VO = 256; +constexpr auto USE_LOGITS_SOFT_CAP = true; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = true; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + DTypeQ* q; + paged_kv_t paged_kv; + DTypeO* o; + float* lse; + + float* maybe_alibi_slopes; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; + + + uint32_t padded_batch_size; + uint32_t num_qo_heads; + IdType q_stride_n; + IdType q_stride_h; + int32_t window_left; + bool enable_pdl; + + IdType* request_indices; + IdType* kv_tile_indices; + IdType* o_indptr; + IdType* kv_chunk_size_ptr; + bool* block_valid_mask; + bool partition_kv; + + __host__ __device__ __forceinline__ int32_t get_qo_len(int32_t batch_idx) const { return 1; } + + __host__ __device__ __forceinline__ int32_t get_kv_len(int32_t batch_idx) const { + return paged_kv.get_length(batch_idx); + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode_jit_pybind.cu b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..91e944770b8cee57481f0631545cef02694f50cc --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode_jit_pybind.cu @@ -0,0 +1,40 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "batch_decode_config.inc" +#include "pytorch_extension_utils.h" + +at::Tensor BatchDecodeWithPagedKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size, + int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, + int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo, + at::Tensor empty_q_data, at::Tensor empty_kv_data); + +void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, + at::Tensor paged_v_cache, at::Tensor paged_kv_indptr, + at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, + int64_t kv_layout_code, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Batched decode with paged KV-Cache plan + m.def("plan", BatchDecodeWithPagedKVCachePlan); + // Batched decode with paged KV-Cache run + m.def("run", BatchDecodeWithPagedKVCacheRun); +} diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode_kernel.cu b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode_kernel.cu new file mode 100644 index 0000000000000000000000000000000000000000..fc08143c5b0cc6d1ba272d41bae81db75865298c --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/batch_decode_kernel.cu @@ -0,0 +1,13 @@ +#include +#include "batch_decode_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t +BatchDecodeWithPagedKVCacheDispatched<256, PosEncodingMode::kNone, DefaultAttention, Params>( + Params params, half* tmp_v, + float* tmp_s, bool enable_pdl, cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode.cu b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode.cu new file mode 100644 index 0000000000000000000000000000000000000000..02c75e02718f8f8ce9b6036f8f788d9a9e254cf1 --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode.cu @@ -0,0 +1,197 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include +#include + +#include "batch_decode_config.inc" +#include "pytorch_conversion_utils.h" +#include "pytorch_extension_utils.h" + +namespace flashinfer { + +template +cudaError_t BatchDecodeWithPagedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v, + float* tmp_s, bool enable_pdl, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +at::Tensor BatchDecodeWithPagedKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size, + int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, + int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo, + at::Tensor empty_q_data, at::Tensor empty_kv_data) { + size_t float_workspace_size_in_bytes = + float_workspace_buffer.size(0) * float_workspace_buffer.element_size(); + size_t int_workspace_size_in_bytes = + int_workspace_buffer.size(0) * int_workspace_buffer.element_size(); + + DecodePlanInfo plan_info; + + auto q_scalar_type = empty_q_data.scalar_type(); + auto kv_scalar_type = empty_kv_data.scalar_type(); + + TORCH_CHECK(head_dim_qk == head_dim_vo, + "CUDA cores template only supports equal head dim for QK and VO, please use tensor " + "cores template for different head dim"); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + DISPATCH_GQA_GROUP_SIZE(num_qo_heads / num_kv_heads, GROUP_SIZE, { + auto work_estimation_func = BatchDecodeWithPagedKVCacheWorkEstimationDispatched< + GROUP_SIZE, HEAD_DIM_QK, POS_ENCODING_MODE, AttentionVariant, Params>; + cudaError_t status = DecodePlan( + static_cast(float_workspace_buffer.data_ptr()), float_workspace_size_in_bytes, + static_cast(int_workspace_buffer.data_ptr()), + static_cast(page_locked_int_workspace_buffer.data_ptr()), + int_workspace_size_in_bytes, plan_info, static_cast(indptr.data_ptr()), + batch_size, num_qo_heads, page_size, enable_cuda_graph, + /*stream=*/stream, work_estimation_func); + + TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); + }); + + return vec_to_tensor(plan_info.ToVector()); +} + +void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, + at::Tensor paged_v_cache, at::Tensor paged_kv_indptr, + at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, + int64_t kv_layout_code, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + DecodePlanInfo plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + QKVLayout kv_layout = static_cast(kv_layout_code); + auto device = q.device(); + int64_t batch_size = q.size(0); + int64_t num_qo_heads = q.size(1); + int64_t num_kv_heads, page_size; + + if (kv_layout == QKVLayout::kHND) { + num_kv_heads = paged_k_cache.size(1); + page_size = paged_k_cache.size(2); + } else { + page_size = paged_k_cache.size(1); + num_kv_heads = paged_k_cache.size(2); + } + uint32_t head_dim_qk = q.size(2); + uint32_t head_dim_vo = paged_v_cache.size(3); + + TORCH_CHECK(head_dim_qk == head_dim_vo, + "CUDA cores template only supports equal head dim for QK and VO, please use tensor " + "cores template for different head dim"); + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == batch_size, lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == num_qo_heads, lse.size(1), q.size(1)); + } + + void* float_buffer = static_cast(float_workspace_buffer.data_ptr()); + void* int_buffer = static_cast(int_workspace_buffer.data_ptr()); + + // get q_scalar_type and kv_scalar_type + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = paged_k_cache.scalar_type(); + + // get q_stride_n and q_stride_h + const auto q_stride_n = q.stride(0); + const auto q_stride_h = q.stride(1); + + // get kv_cache_strides + const int64_t* kv_cache_strides = nullptr; + auto k_strides = paged_k_cache.strides(); + auto v_strides = paged_v_cache.strides(); + TORCH_CHECK(k_strides == v_strides, "k/v strides must be identical"); + kv_cache_strides = k_strides.data(); + + const c10::cuda::OptionalCUDAGuard device_guard(device); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + paged_kv_t paged_kv( + num_kv_heads, page_size, HEAD_DIM_QK, batch_size, kv_layout, + static_cast(paged_k_cache.data_ptr()), + static_cast(paged_v_cache.data_ptr()), kv_cache_strides, + static_cast(paged_kv_indices.data_ptr()), + static_cast(paged_kv_indptr.data_ptr()), + static_cast(paged_kv_last_page_len.data_ptr())); + + Params params; + params.q = static_cast(q.data_ptr()); + params.paged_kv = paged_kv; + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.padded_batch_size = 0; + params.num_qo_heads = num_qo_heads; + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.window_left = window_left; + params.request_indices = nullptr; + params.kv_tile_indices = nullptr; + params.o_indptr = nullptr; + params.kv_chunk_size_ptr = nullptr; + params.block_valid_mask = nullptr; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + DTypeO* tmp_v = nullptr; + float* tmp_s = nullptr; + params.request_indices = + GetPtrFromBaseOffset(int_buffer, plan_info.request_indices_offset); + params.kv_tile_indices = + GetPtrFromBaseOffset(int_buffer, plan_info.kv_tile_indices_offset); + params.o_indptr = GetPtrFromBaseOffset(int_buffer, plan_info.o_indptr_offset); + params.kv_chunk_size_ptr = + GetPtrFromBaseOffset(int_buffer, plan_info.kv_chunk_size_ptr_offset); + if (plan_info.split_kv) { + tmp_v = GetPtrFromBaseOffset(float_buffer, plan_info.v_offset); + tmp_s = GetPtrFromBaseOffset(float_buffer, plan_info.s_offset); + if (plan_info.enable_cuda_graph) { + params.block_valid_mask = + GetPtrFromBaseOffset(int_buffer, plan_info.block_valid_mask_offset); + } + } + params.padded_batch_size = plan_info.padded_batch_size; + + cudaError_t status = + flashinfer::BatchDecodeWithPagedKVCacheDispatched(params, tmp_v, + tmp_s, enable_pdl, + /*stream=*/stream); + TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); +} diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_config.inc b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..f64f4467866940f9acd1a43d340583489bdf05f6 --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_config.inc @@ -0,0 +1,71 @@ +#pragma once +#include +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta +#define ADDITIONAL_PARAMS_SETTER params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, ...) { \ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ +} + +using namespace flashinfer; + +using DTypeQ = half; +using DTypeKV = half; +using DTypeO = half; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 64; +constexpr int HEAD_DIM_VO = 64; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = false; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + DTypeQ* q; + paged_kv_t paged_kv; + DTypeO* o; + float* lse; + + float* maybe_alibi_slopes; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; + + + uint32_t padded_batch_size; + uint32_t num_qo_heads; + IdType q_stride_n; + IdType q_stride_h; + int32_t window_left; + bool enable_pdl; + + IdType* request_indices; + IdType* kv_tile_indices; + IdType* o_indptr; + IdType* kv_chunk_size_ptr; + bool* block_valid_mask; + bool partition_kv; + + __host__ __device__ __forceinline__ int32_t get_qo_len(int32_t batch_idx) const { return 1; } + + __host__ __device__ __forceinline__ int32_t get_kv_len(int32_t batch_idx) const { + return paged_kv.get_length(batch_idx); + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_jit_pybind.cu b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..91e944770b8cee57481f0631545cef02694f50cc --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_jit_pybind.cu @@ -0,0 +1,40 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "batch_decode_config.inc" +#include "pytorch_extension_utils.h" + +at::Tensor BatchDecodeWithPagedKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size, + int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, + int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo, + at::Tensor empty_q_data, at::Tensor empty_kv_data); + +void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, + at::Tensor paged_v_cache, at::Tensor paged_kv_indptr, + at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, + int64_t kv_layout_code, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Batched decode with paged KV-Cache plan + m.def("plan", BatchDecodeWithPagedKVCachePlan); + // Batched decode with paged KV-Cache run + m.def("run", BatchDecodeWithPagedKVCacheRun); +} diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_kernel.cu b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_kernel.cu new file mode 100644 index 0000000000000000000000000000000000000000..91336af2f35e622c26c8ef0518429688a1839b6a --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/batch_decode_kernel.cu @@ -0,0 +1,13 @@ +#include +#include "batch_decode_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t +BatchDecodeWithPagedKVCacheDispatched<64, PosEncodingMode::kNone, DefaultAttention, Params>( + Params params, half* tmp_v, + float* tmp_s, bool enable_pdl, cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode.cu b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode.cu new file mode 100644 index 0000000000000000000000000000000000000000..02c75e02718f8f8ce9b6036f8f788d9a9e254cf1 --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode.cu @@ -0,0 +1,197 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include +#include + +#include "batch_decode_config.inc" +#include "pytorch_conversion_utils.h" +#include "pytorch_extension_utils.h" + +namespace flashinfer { + +template +cudaError_t BatchDecodeWithPagedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v, + float* tmp_s, bool enable_pdl, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +at::Tensor BatchDecodeWithPagedKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size, + int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, + int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo, + at::Tensor empty_q_data, at::Tensor empty_kv_data) { + size_t float_workspace_size_in_bytes = + float_workspace_buffer.size(0) * float_workspace_buffer.element_size(); + size_t int_workspace_size_in_bytes = + int_workspace_buffer.size(0) * int_workspace_buffer.element_size(); + + DecodePlanInfo plan_info; + + auto q_scalar_type = empty_q_data.scalar_type(); + auto kv_scalar_type = empty_kv_data.scalar_type(); + + TORCH_CHECK(head_dim_qk == head_dim_vo, + "CUDA cores template only supports equal head dim for QK and VO, please use tensor " + "cores template for different head dim"); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + DISPATCH_GQA_GROUP_SIZE(num_qo_heads / num_kv_heads, GROUP_SIZE, { + auto work_estimation_func = BatchDecodeWithPagedKVCacheWorkEstimationDispatched< + GROUP_SIZE, HEAD_DIM_QK, POS_ENCODING_MODE, AttentionVariant, Params>; + cudaError_t status = DecodePlan( + static_cast(float_workspace_buffer.data_ptr()), float_workspace_size_in_bytes, + static_cast(int_workspace_buffer.data_ptr()), + static_cast(page_locked_int_workspace_buffer.data_ptr()), + int_workspace_size_in_bytes, plan_info, static_cast(indptr.data_ptr()), + batch_size, num_qo_heads, page_size, enable_cuda_graph, + /*stream=*/stream, work_estimation_func); + + TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); + }); + + return vec_to_tensor(plan_info.ToVector()); +} + +void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, + at::Tensor paged_v_cache, at::Tensor paged_kv_indptr, + at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, + int64_t kv_layout_code, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + DecodePlanInfo plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + QKVLayout kv_layout = static_cast(kv_layout_code); + auto device = q.device(); + int64_t batch_size = q.size(0); + int64_t num_qo_heads = q.size(1); + int64_t num_kv_heads, page_size; + + if (kv_layout == QKVLayout::kHND) { + num_kv_heads = paged_k_cache.size(1); + page_size = paged_k_cache.size(2); + } else { + page_size = paged_k_cache.size(1); + num_kv_heads = paged_k_cache.size(2); + } + uint32_t head_dim_qk = q.size(2); + uint32_t head_dim_vo = paged_v_cache.size(3); + + TORCH_CHECK(head_dim_qk == head_dim_vo, + "CUDA cores template only supports equal head dim for QK and VO, please use tensor " + "cores template for different head dim"); + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == batch_size, lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == num_qo_heads, lse.size(1), q.size(1)); + } + + void* float_buffer = static_cast(float_workspace_buffer.data_ptr()); + void* int_buffer = static_cast(int_workspace_buffer.data_ptr()); + + // get q_scalar_type and kv_scalar_type + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = paged_k_cache.scalar_type(); + + // get q_stride_n and q_stride_h + const auto q_stride_n = q.stride(0); + const auto q_stride_h = q.stride(1); + + // get kv_cache_strides + const int64_t* kv_cache_strides = nullptr; + auto k_strides = paged_k_cache.strides(); + auto v_strides = paged_v_cache.strides(); + TORCH_CHECK(k_strides == v_strides, "k/v strides must be identical"); + kv_cache_strides = k_strides.data(); + + const c10::cuda::OptionalCUDAGuard device_guard(device); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + paged_kv_t paged_kv( + num_kv_heads, page_size, HEAD_DIM_QK, batch_size, kv_layout, + static_cast(paged_k_cache.data_ptr()), + static_cast(paged_v_cache.data_ptr()), kv_cache_strides, + static_cast(paged_kv_indices.data_ptr()), + static_cast(paged_kv_indptr.data_ptr()), + static_cast(paged_kv_last_page_len.data_ptr())); + + Params params; + params.q = static_cast(q.data_ptr()); + params.paged_kv = paged_kv; + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.padded_batch_size = 0; + params.num_qo_heads = num_qo_heads; + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.window_left = window_left; + params.request_indices = nullptr; + params.kv_tile_indices = nullptr; + params.o_indptr = nullptr; + params.kv_chunk_size_ptr = nullptr; + params.block_valid_mask = nullptr; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + DTypeO* tmp_v = nullptr; + float* tmp_s = nullptr; + params.request_indices = + GetPtrFromBaseOffset(int_buffer, plan_info.request_indices_offset); + params.kv_tile_indices = + GetPtrFromBaseOffset(int_buffer, plan_info.kv_tile_indices_offset); + params.o_indptr = GetPtrFromBaseOffset(int_buffer, plan_info.o_indptr_offset); + params.kv_chunk_size_ptr = + GetPtrFromBaseOffset(int_buffer, plan_info.kv_chunk_size_ptr_offset); + if (plan_info.split_kv) { + tmp_v = GetPtrFromBaseOffset(float_buffer, plan_info.v_offset); + tmp_s = GetPtrFromBaseOffset(float_buffer, plan_info.s_offset); + if (plan_info.enable_cuda_graph) { + params.block_valid_mask = + GetPtrFromBaseOffset(int_buffer, plan_info.block_valid_mask_offset); + } + } + params.padded_batch_size = plan_info.padded_batch_size; + + cudaError_t status = + flashinfer::BatchDecodeWithPagedKVCacheDispatched(params, tmp_v, + tmp_s, enable_pdl, + /*stream=*/stream); + TORCH_CHECK(status == cudaSuccess, "BatchDecodeWithPagedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); +} diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode_config.inc b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..6a2d66570d484554059fda2a9d86902b9995c1f0 --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode_config.inc @@ -0,0 +1,71 @@ +#pragma once +#include +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta +#define ADDITIONAL_PARAMS_SETTER params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, ...) { \ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ +} + +using namespace flashinfer; + +using DTypeQ = half; +using DTypeKV = half; +using DTypeO = half; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 64; +constexpr int HEAD_DIM_VO = 64; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = true; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + DTypeQ* q; + paged_kv_t paged_kv; + DTypeO* o; + float* lse; + + float* maybe_alibi_slopes; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; + + + uint32_t padded_batch_size; + uint32_t num_qo_heads; + IdType q_stride_n; + IdType q_stride_h; + int32_t window_left; + bool enable_pdl; + + IdType* request_indices; + IdType* kv_tile_indices; + IdType* o_indptr; + IdType* kv_chunk_size_ptr; + bool* block_valid_mask; + bool partition_kv; + + __host__ __device__ __forceinline__ int32_t get_qo_len(int32_t batch_idx) const { return 1; } + + __host__ __device__ __forceinline__ int32_t get_kv_len(int32_t batch_idx) const { + return paged_kv.get_length(batch_idx); + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode_jit_pybind.cu b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..91e944770b8cee57481f0631545cef02694f50cc --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode_jit_pybind.cu @@ -0,0 +1,40 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "batch_decode_config.inc" +#include "pytorch_extension_utils.h" + +at::Tensor BatchDecodeWithPagedKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor indptr, int64_t batch_size, + int64_t num_qo_heads, int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, + int64_t window_left, double logits_soft_cap, int64_t head_dim_qk, int64_t head_dim_vo, + at::Tensor empty_q_data, at::Tensor empty_kv_data); + +void BatchDecodeWithPagedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, + at::Tensor paged_v_cache, at::Tensor paged_kv_indptr, + at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, + int64_t kv_layout_code, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Batched decode with paged KV-Cache plan + m.def("plan", BatchDecodeWithPagedKVCachePlan); + // Batched decode with paged KV-Cache run + m.def("run", BatchDecodeWithPagedKVCacheRun); +} diff --git a/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode_kernel.cu b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode_kernel.cu new file mode 100644 index 0000000000000000000000000000000000000000..f9cecd66f3d2b259979d5fb2b612e9212bb3b57a --- /dev/null +++ b/csrc/generated/batch_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/batch_decode_kernel.cu @@ -0,0 +1,13 @@ +#include +#include "batch_decode_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t +BatchDecodeWithPagedKVCacheDispatched<64, PosEncodingMode::kNone, DefaultAttention, Params>( + Params params, half* tmp_v, + float* tmp_s, bool enable_pdl, cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/batch_mla_attention_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_ckv_512_head_dim_kpe_64_profiler_False/batch_mla_config.inc b/csrc/generated/batch_mla_attention_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_ckv_512_head_dim_kpe_64_profiler_False/batch_mla_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..c9dcaa1614d2ac92a527ce44b85a13f286cd251c --- /dev/null +++ b/csrc/generated/batch_mla_attention_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_ckv_512_head_dim_kpe_64_profiler_False/batch_mla_config.inc @@ -0,0 +1,33 @@ +#pragma once +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace flashinfer; + +#ifdef FLASHINFER_ENABLE_PROFILER +#define ADDITIONAL_FUNC_PARAMS , at::Tensor profiler_buffer +#define ADDITIONAL_PARAMS_SETTER \ + params.profiler_buffer = static_cast(profiler_buffer.data_ptr()); +#else +#define ADDITIONAL_FUNC_PARAMS +#define ADDITIONAL_PARAMS_SETTER +#endif + +using DTypeQ = nv_bfloat16; +using DTypeKV = nv_bfloat16; +using DTypeO = nv_bfloat16; +using IdType = int32_t; +constexpr int HEAD_DIM_CKV = 512; +constexpr int HEAD_DIM_KPE = 64; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_CKV, HEAD_DIM_KPE, Params, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { \ + using Params = MLAParams; \ + __VA_ARGS__(); \ + }) \ No newline at end of file diff --git a/csrc/generated/batch_mla_attention_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_ckv_512_head_dim_kpe_64_profiler_False/batch_mla_plan.cu b/csrc/generated/batch_mla_attention_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_ckv_512_head_dim_kpe_64_profiler_False/batch_mla_plan.cu new file mode 100644 index 0000000000000000000000000000000000000000..0276102cab54e4ed1fe225526b4e8c94b50b5d12 --- /dev/null +++ b/csrc/generated/batch_mla_attention_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_ckv_512_head_dim_kpe_64_profiler_False/batch_mla_plan.cu @@ -0,0 +1,52 @@ +/* + * Copyright (c) 2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include + +#include "batch_mla_config.inc" +#include "pytorch_conversion_utils.h" +#include "pytorch_extension_utils.h" + +using namespace flashinfer; + +at::Tensor BatchMLAPagedAttentionPlan(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, + at::Tensor qo_indptr, at::Tensor kv_indptr, at::Tensor kv_len, + int64_t num_heads, int64_t head_dim_o, bool causal) { + size_t float_workspace_size_in_bytes = + float_workspace_buffer.size(0) * float_workspace_buffer.element_size(); + size_t int_workspace_size_in_bytes = + int_workspace_buffer.size(0) * int_workspace_buffer.element_size(); + + MLAPlanInfo plan_info; + + int batch_size = kv_len.size(0); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + cudaError_t status = + MLAPlan(float_workspace_buffer.data_ptr(), float_workspace_size_in_bytes, + int_workspace_buffer.data_ptr(), page_locked_int_workspace_buffer.data_ptr(), + int_workspace_size_in_bytes, plan_info, static_cast(qo_indptr.data_ptr()), + static_cast(kv_indptr.data_ptr()), static_cast(kv_len.data_ptr()), + batch_size, num_heads, head_dim_o, causal, stream); + + TORCH_CHECK(status == cudaSuccess, "Failed to plan MLA, error: ", cudaGetErrorString(status)); + + return vec_to_tensor(plan_info.ToVector()); +} diff --git a/csrc/generated/batch_mla_attention_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_ckv_512_head_dim_kpe_64_profiler_False/batch_mla_pybind.cu b/csrc/generated/batch_mla_attention_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_ckv_512_head_dim_kpe_64_profiler_False/batch_mla_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..438fc39519ddf51cc40337af1eeccd2c88163c31 --- /dev/null +++ b/csrc/generated/batch_mla_attention_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_ckv_512_head_dim_kpe_64_profiler_False/batch_mla_pybind.cu @@ -0,0 +1,35 @@ +/* + * Copyright (c) 2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "batch_mla_config.inc" +#include "pytorch_extension_utils.h" + +at::Tensor BatchMLAPagedAttentionPlan(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, + at::Tensor qo_indptr, at::Tensor kv_indptr, at::Tensor kv_len, + int64_t num_heads, int64_t head_dim_o, bool causal); + +void BatchMLAPagedAttentionRun(at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor plan_info_vec, at::Tensor q_nope, at::Tensor q_pe, + at::Tensor ckv_cache, at::Tensor kpe_cache, at::Tensor kv_indices, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t num_heads, int64_t page_size, + double sm_scale); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + m.def("plan", &BatchMLAPagedAttentionPlan); + m.def("run", &BatchMLAPagedAttentionRun); +} diff --git a/csrc/generated/batch_mla_attention_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_ckv_512_head_dim_kpe_64_profiler_False/batch_mla_run.cu b/csrc/generated/batch_mla_attention_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_ckv_512_head_dim_kpe_64_profiler_False/batch_mla_run.cu new file mode 100644 index 0000000000000000000000000000000000000000..ef4c4848ea822a7c54e5b980823865489c50e18e --- /dev/null +++ b/csrc/generated/batch_mla_attention_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_ckv_512_head_dim_kpe_64_profiler_False/batch_mla_run.cu @@ -0,0 +1,125 @@ +/* + * Copyright (c) 2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include +#include + +#include "batch_mla_config.inc" +#include "pytorch_conversion_utils.h" +#include "pytorch_extension_utils.h" + +using namespace flashinfer; + +void BatchMLAPagedAttentionRun(at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor plan_info_vec, at::Tensor q_nope, at::Tensor q_pe, + at::Tensor ckv_cache, at::Tensor kpe_cache, at::Tensor kv_indices, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t num_heads, int64_t page_size, + double sm_scale) { + // q_nope: [n, num_heads, head_dim_ckv] + // q_pe: [n, num_heads, head_dim_kpe] + // ckv_cache: [num_pages, page_size, head_dim_ckv] + // kpe_cache: [num_pages, page_size, head_dim_kpe] + MLAPlanInfo plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + + auto device = q_nope.device(); + + void* float_buffer_ptr = float_workspace_buffer.data_ptr(); + void* int_buffer_ptr = int_workspace_buffer.data_ptr(); + + const MaskMode mask_mode = static_cast(mask_mode_code); + + auto q_scalar_type = q_nope.scalar_type(); + auto kv_scalar_type = ckv_cache.scalar_type(); + + unsigned int q_nope_stride_n = q_nope.stride(0); + unsigned int q_nope_stride_h = q_nope.stride(1); + unsigned int q_pe_stride_n = q_pe.stride(0); + unsigned int q_pe_stride_h = q_pe.stride(1); + unsigned int ckv_stride_page = ckv_cache.stride(0); + unsigned int ckv_stride_n = ckv_cache.stride(1); + unsigned int kpe_stride_page = kpe_cache.stride(0); + unsigned int kpe_stride_n = kpe_cache.stride(1); + unsigned int o_stride_n = o.stride(0); + unsigned int o_stride_h = o.stride(1); + + const c10::cuda::OptionalCUDAGuard device_guard(device); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_CKV, HEAD_DIM_KPE, Params, [&] { + Params params; + + params.q_nope = static_cast(q_nope.data_ptr()); + params.q_pe = static_cast(q_pe.data_ptr()); + params.ckv = static_cast(ckv_cache.data_ptr()); + params.kpe = static_cast(kpe_cache.data_ptr()); + + params.q_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.q_indptr_offset); + params.kv_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_indptr_offset); + params.partial_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.partial_indptr_offset); + params.kv_indices = static_cast(kv_indices.data_ptr()); + params.q_len = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.q_len_offset); + params.kv_len = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_len_offset); + params.q_start = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.q_start_offset); + params.kv_start = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_start_offset); + params.kv_end = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_end_offset); + params.work_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.work_indptr_offset); + params.merge_packed_offset_start = GetPtrFromBaseOffset( + int_buffer_ptr, plan_info.merge_packed_offset_start_offset); + params.merge_packed_offset_end = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.merge_packed_offset_end_offset); + params.merge_partial_packed_offset_start = GetPtrFromBaseOffset( + int_buffer_ptr, plan_info.merge_partial_packed_offset_start_offset); + params.merge_partial_packed_offset_end = GetPtrFromBaseOffset( + int_buffer_ptr, plan_info.merge_partial_packed_offset_end_offset); + params.merge_partial_stride = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.merge_partial_stride_offset); + params.final_o = static_cast(o.data_ptr()); + params.final_lse = + maybe_lse.has_value() ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.partial_o = + GetPtrFromBaseOffset(float_buffer_ptr, plan_info.partial_o_offset); + params.partial_lse = + GetPtrFromBaseOffset(float_buffer_ptr, plan_info.partial_lse_offset); + + params.num_heads = uint_fastdiv(num_heads); + params.block_size = uint_fastdiv(page_size); + + params.q_nope_stride_n = q_nope_stride_n; + params.q_nope_stride_h = q_nope_stride_h; + params.q_pe_stride_n = q_pe_stride_n; + params.q_pe_stride_h = q_pe_stride_h; + params.ckv_stride_page = ckv_stride_page; + params.ckv_stride_n = ckv_stride_n; + params.kpe_stride_page = kpe_stride_page; + params.kpe_stride_n = kpe_stride_n; + params.o_stride_n = o_stride_n; + params.o_stride_h = o_stride_h; + + params.sm_scale = sm_scale; + + cudaError_t status = mla::BatchMLAPagedAttention( + params, plan_info.num_blks_x, plan_info.num_blks_y, stream); + + TORCH_CHECK(status == cudaSuccess, + "Failed to run MLA, error: ", cudaGetErrorString(status)); + }); +} diff --git a/csrc/generated/batch_mla_attention_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_ckv_512_head_dim_kpe_64_profiler_False_sm90/batch_mla_sm90_config.inc b/csrc/generated/batch_mla_attention_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_ckv_512_head_dim_kpe_64_profiler_False_sm90/batch_mla_sm90_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..c9dcaa1614d2ac92a527ce44b85a13f286cd251c --- /dev/null +++ b/csrc/generated/batch_mla_attention_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_ckv_512_head_dim_kpe_64_profiler_False_sm90/batch_mla_sm90_config.inc @@ -0,0 +1,33 @@ +#pragma once +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace flashinfer; + +#ifdef FLASHINFER_ENABLE_PROFILER +#define ADDITIONAL_FUNC_PARAMS , at::Tensor profiler_buffer +#define ADDITIONAL_PARAMS_SETTER \ + params.profiler_buffer = static_cast(profiler_buffer.data_ptr()); +#else +#define ADDITIONAL_FUNC_PARAMS +#define ADDITIONAL_PARAMS_SETTER +#endif + +using DTypeQ = nv_bfloat16; +using DTypeKV = nv_bfloat16; +using DTypeO = nv_bfloat16; +using IdType = int32_t; +constexpr int HEAD_DIM_CKV = 512; +constexpr int HEAD_DIM_KPE = 64; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_CKV, HEAD_DIM_KPE, Params, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { \ + using Params = MLAParams; \ + __VA_ARGS__(); \ + }) \ No newline at end of file diff --git a/csrc/generated/batch_mla_attention_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_ckv_512_head_dim_kpe_64_profiler_False_sm90/batch_mla_sm90_plan.cu b/csrc/generated/batch_mla_attention_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_ckv_512_head_dim_kpe_64_profiler_False_sm90/batch_mla_sm90_plan.cu new file mode 100644 index 0000000000000000000000000000000000000000..35f6d0fa820ba335fdb04324c16f054ab66d2279 --- /dev/null +++ b/csrc/generated/batch_mla_attention_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_ckv_512_head_dim_kpe_64_profiler_False_sm90/batch_mla_sm90_plan.cu @@ -0,0 +1,53 @@ +/* + * Copyright (c) 2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include + +#include "batch_mla_sm90_config.inc" +#include "pytorch_conversion_utils.h" +#include "pytorch_extension_utils.h" + +using namespace flashinfer; + +at::Tensor BatchMLAPagedAttentionSM90Plan(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, + at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len, int64_t num_heads, int64_t head_dim_o, + bool causal) { + size_t float_workspace_size_in_bytes = + float_workspace_buffer.size(0) * float_workspace_buffer.element_size(); + size_t int_workspace_size_in_bytes = + int_workspace_buffer.size(0) * int_workspace_buffer.element_size(); + + MLAPlanInfo plan_info; + + int batch_size = kv_len.size(0); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + cudaError_t status = + MLAPlan(float_workspace_buffer.data_ptr(), float_workspace_size_in_bytes, + int_workspace_buffer.data_ptr(), page_locked_int_workspace_buffer.data_ptr(), + int_workspace_size_in_bytes, plan_info, static_cast(qo_indptr.data_ptr()), + static_cast(kv_indptr.data_ptr()), static_cast(kv_len.data_ptr()), + batch_size, num_heads, head_dim_o, causal, stream); + + TORCH_CHECK(status == cudaSuccess, "Failed to plan MLA, error: ", cudaGetErrorString(status)); + + return vec_to_tensor(plan_info.ToVector()); +} diff --git a/csrc/generated/batch_mla_attention_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_ckv_512_head_dim_kpe_64_profiler_False_sm90/batch_mla_sm90_pybind.cu b/csrc/generated/batch_mla_attention_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_ckv_512_head_dim_kpe_64_profiler_False_sm90/batch_mla_sm90_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..d089ad0c5c1acbd1c3ed788b8412d1ae0523ca49 --- /dev/null +++ b/csrc/generated/batch_mla_attention_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_ckv_512_head_dim_kpe_64_profiler_False_sm90/batch_mla_sm90_pybind.cu @@ -0,0 +1,37 @@ +/* + * Copyright (c) 2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "batch_mla_sm90_config.inc" +#include "pytorch_extension_utils.h" + +at::Tensor BatchMLAPagedAttentionSM90Plan(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, + at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len, int64_t num_heads, int64_t head_dim_o, + bool causal); + +void BatchMLAPagedAttentionSM90Run(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q_nope, at::Tensor q_pe, at::Tensor ckv_cache, + at::Tensor kpe_cache, at::Tensor kv_indices, at::Tensor o, + std::optional maybe_lse, int64_t mask_mode_code, + int64_t num_heads, int64_t page_size, + double sm_scale ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + m.def("plan", &BatchMLAPagedAttentionSM90Plan); + m.def("run", &BatchMLAPagedAttentionSM90Run); +} diff --git a/csrc/generated/batch_mla_attention_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_ckv_512_head_dim_kpe_64_profiler_False_sm90/batch_mla_sm90_run.cu b/csrc/generated/batch_mla_attention_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_ckv_512_head_dim_kpe_64_profiler_False_sm90/batch_mla_sm90_run.cu new file mode 100644 index 0000000000000000000000000000000000000000..3da072eb039d16a3cf7ee319be03c487ac6fc288 --- /dev/null +++ b/csrc/generated/batch_mla_attention_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_ckv_512_head_dim_kpe_64_profiler_False_sm90/batch_mla_sm90_run.cu @@ -0,0 +1,129 @@ +/* + * Copyright (c) 2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include +#include + +#include "batch_mla_sm90_config.inc" +#include "pytorch_conversion_utils.h" +#include "pytorch_extension_utils.h" + +using namespace flashinfer; + +void BatchMLAPagedAttentionSM90Run(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q_nope, at::Tensor q_pe, at::Tensor ckv_cache, + at::Tensor kpe_cache, at::Tensor kv_indices, at::Tensor o, + std::optional maybe_lse, int64_t mask_mode_code, + int64_t num_heads, int64_t page_size, + double sm_scale ADDITIONAL_FUNC_PARAMS) { + // q_nope: [n, num_heads, head_dim_ckv] + // q_pe: [n, num_heads, head_dim_kpe] + // ckv_cache: [num_pages, page_size, head_dim_ckv] + // kpe_cache: [num_pages, page_size, head_dim_kpe] + MLAPlanInfo plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + + auto device = q_nope.device(); + + void* float_buffer_ptr = float_workspace_buffer.data_ptr(); + void* int_buffer_ptr = int_workspace_buffer.data_ptr(); + + const MaskMode mask_mode = static_cast(mask_mode_code); + + auto q_scalar_type = q_nope.scalar_type(); + auto kv_scalar_type = ckv_cache.scalar_type(); + + unsigned int q_nope_stride_n = q_nope.stride(0); + unsigned int q_nope_stride_h = q_nope.stride(1); + unsigned int q_pe_stride_n = q_pe.stride(0); + unsigned int q_pe_stride_h = q_pe.stride(1); + unsigned int ckv_stride_page = ckv_cache.stride(0); + unsigned int ckv_stride_n = ckv_cache.stride(1); + unsigned int kpe_stride_page = kpe_cache.stride(0); + unsigned int kpe_stride_n = kpe_cache.stride(1); + unsigned int o_stride_n = o.stride(0); + unsigned int o_stride_h = o.stride(1); + + const c10::cuda::OptionalCUDAGuard device_guard(device); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_CKV, HEAD_DIM_KPE, Params, [&] { + Params params; + + params.q_nope = static_cast(q_nope.data_ptr()); + params.q_pe = static_cast(q_pe.data_ptr()); + params.ckv = static_cast(ckv_cache.data_ptr()); + params.kpe = static_cast(kpe_cache.data_ptr()); + + params.q_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.q_indptr_offset); + params.kv_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_indptr_offset); + params.partial_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.partial_indptr_offset); + params.kv_indices = static_cast(kv_indices.data_ptr()); + params.q_len = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.q_len_offset); + params.kv_len = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_len_offset); + params.q_start = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.q_start_offset); + params.kv_start = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_start_offset); + params.kv_end = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_end_offset); + params.work_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.work_indptr_offset); + params.merge_packed_offset_start = GetPtrFromBaseOffset( + int_buffer_ptr, plan_info.merge_packed_offset_start_offset); + params.merge_packed_offset_end = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.merge_packed_offset_end_offset); + params.merge_partial_packed_offset_start = GetPtrFromBaseOffset( + int_buffer_ptr, plan_info.merge_partial_packed_offset_start_offset); + params.merge_partial_packed_offset_end = GetPtrFromBaseOffset( + int_buffer_ptr, plan_info.merge_partial_packed_offset_end_offset); + params.merge_partial_stride = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.merge_partial_stride_offset); + params.final_o = static_cast(o.data_ptr()); + params.final_lse = + maybe_lse.has_value() ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.partial_o = + GetPtrFromBaseOffset(float_buffer_ptr, plan_info.partial_o_offset); + params.partial_lse = + GetPtrFromBaseOffset(float_buffer_ptr, plan_info.partial_lse_offset); + + params.num_heads = uint_fastdiv(num_heads); + params.block_size = uint_fastdiv(page_size); + + params.q_nope_stride_n = q_nope_stride_n; + params.q_nope_stride_h = q_nope_stride_h; + params.q_pe_stride_n = q_pe_stride_n; + params.q_pe_stride_h = q_pe_stride_h; + params.ckv_stride_page = ckv_stride_page; + params.ckv_stride_n = ckv_stride_n; + params.kpe_stride_page = kpe_stride_page; + params.kpe_stride_n = kpe_stride_n; + params.o_stride_n = o_stride_n; + params.o_stride_h = o_stride_h; + + ADDITIONAL_PARAMS_SETTER + + params.sm_scale = sm_scale; + + cudaError_t status = + mla::BatchMLAPageAttentionHopper( + params, plan_info.num_blks_x, plan_info.num_blks_y, stream); + + TORCH_CHECK(status == cudaSuccess, + "Failed to run MLA, error: ", cudaGetErrorString(status)); + }); +} diff --git a/csrc/generated/batch_mla_attention_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_ckv_512_head_dim_kpe_64_profiler_False/batch_mla_config.inc b/csrc/generated/batch_mla_attention_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_ckv_512_head_dim_kpe_64_profiler_False/batch_mla_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..6f4c28dc8ea7f7ab624ed810abb6ee6fe7cfccb7 --- /dev/null +++ b/csrc/generated/batch_mla_attention_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_ckv_512_head_dim_kpe_64_profiler_False/batch_mla_config.inc @@ -0,0 +1,33 @@ +#pragma once +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace flashinfer; + +#ifdef FLASHINFER_ENABLE_PROFILER +#define ADDITIONAL_FUNC_PARAMS , at::Tensor profiler_buffer +#define ADDITIONAL_PARAMS_SETTER \ + params.profiler_buffer = static_cast(profiler_buffer.data_ptr()); +#else +#define ADDITIONAL_FUNC_PARAMS +#define ADDITIONAL_PARAMS_SETTER +#endif + +using DTypeQ = half; +using DTypeKV = half; +using DTypeO = half; +using IdType = int32_t; +constexpr int HEAD_DIM_CKV = 512; +constexpr int HEAD_DIM_KPE = 64; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_CKV, HEAD_DIM_KPE, Params, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { \ + using Params = MLAParams; \ + __VA_ARGS__(); \ + }) \ No newline at end of file diff --git a/csrc/generated/batch_mla_attention_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_ckv_512_head_dim_kpe_64_profiler_False/batch_mla_plan.cu b/csrc/generated/batch_mla_attention_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_ckv_512_head_dim_kpe_64_profiler_False/batch_mla_plan.cu new file mode 100644 index 0000000000000000000000000000000000000000..0276102cab54e4ed1fe225526b4e8c94b50b5d12 --- /dev/null +++ b/csrc/generated/batch_mla_attention_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_ckv_512_head_dim_kpe_64_profiler_False/batch_mla_plan.cu @@ -0,0 +1,52 @@ +/* + * Copyright (c) 2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include + +#include "batch_mla_config.inc" +#include "pytorch_conversion_utils.h" +#include "pytorch_extension_utils.h" + +using namespace flashinfer; + +at::Tensor BatchMLAPagedAttentionPlan(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, + at::Tensor qo_indptr, at::Tensor kv_indptr, at::Tensor kv_len, + int64_t num_heads, int64_t head_dim_o, bool causal) { + size_t float_workspace_size_in_bytes = + float_workspace_buffer.size(0) * float_workspace_buffer.element_size(); + size_t int_workspace_size_in_bytes = + int_workspace_buffer.size(0) * int_workspace_buffer.element_size(); + + MLAPlanInfo plan_info; + + int batch_size = kv_len.size(0); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + cudaError_t status = + MLAPlan(float_workspace_buffer.data_ptr(), float_workspace_size_in_bytes, + int_workspace_buffer.data_ptr(), page_locked_int_workspace_buffer.data_ptr(), + int_workspace_size_in_bytes, plan_info, static_cast(qo_indptr.data_ptr()), + static_cast(kv_indptr.data_ptr()), static_cast(kv_len.data_ptr()), + batch_size, num_heads, head_dim_o, causal, stream); + + TORCH_CHECK(status == cudaSuccess, "Failed to plan MLA, error: ", cudaGetErrorString(status)); + + return vec_to_tensor(plan_info.ToVector()); +} diff --git a/csrc/generated/batch_mla_attention_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_ckv_512_head_dim_kpe_64_profiler_False/batch_mla_pybind.cu b/csrc/generated/batch_mla_attention_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_ckv_512_head_dim_kpe_64_profiler_False/batch_mla_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..438fc39519ddf51cc40337af1eeccd2c88163c31 --- /dev/null +++ b/csrc/generated/batch_mla_attention_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_ckv_512_head_dim_kpe_64_profiler_False/batch_mla_pybind.cu @@ -0,0 +1,35 @@ +/* + * Copyright (c) 2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "batch_mla_config.inc" +#include "pytorch_extension_utils.h" + +at::Tensor BatchMLAPagedAttentionPlan(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, + at::Tensor qo_indptr, at::Tensor kv_indptr, at::Tensor kv_len, + int64_t num_heads, int64_t head_dim_o, bool causal); + +void BatchMLAPagedAttentionRun(at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor plan_info_vec, at::Tensor q_nope, at::Tensor q_pe, + at::Tensor ckv_cache, at::Tensor kpe_cache, at::Tensor kv_indices, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t num_heads, int64_t page_size, + double sm_scale); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + m.def("plan", &BatchMLAPagedAttentionPlan); + m.def("run", &BatchMLAPagedAttentionRun); +} diff --git a/csrc/generated/batch_mla_attention_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_ckv_512_head_dim_kpe_64_profiler_False/batch_mla_run.cu b/csrc/generated/batch_mla_attention_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_ckv_512_head_dim_kpe_64_profiler_False/batch_mla_run.cu new file mode 100644 index 0000000000000000000000000000000000000000..ef4c4848ea822a7c54e5b980823865489c50e18e --- /dev/null +++ b/csrc/generated/batch_mla_attention_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_ckv_512_head_dim_kpe_64_profiler_False/batch_mla_run.cu @@ -0,0 +1,125 @@ +/* + * Copyright (c) 2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include +#include + +#include "batch_mla_config.inc" +#include "pytorch_conversion_utils.h" +#include "pytorch_extension_utils.h" + +using namespace flashinfer; + +void BatchMLAPagedAttentionRun(at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor plan_info_vec, at::Tensor q_nope, at::Tensor q_pe, + at::Tensor ckv_cache, at::Tensor kpe_cache, at::Tensor kv_indices, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t num_heads, int64_t page_size, + double sm_scale) { + // q_nope: [n, num_heads, head_dim_ckv] + // q_pe: [n, num_heads, head_dim_kpe] + // ckv_cache: [num_pages, page_size, head_dim_ckv] + // kpe_cache: [num_pages, page_size, head_dim_kpe] + MLAPlanInfo plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + + auto device = q_nope.device(); + + void* float_buffer_ptr = float_workspace_buffer.data_ptr(); + void* int_buffer_ptr = int_workspace_buffer.data_ptr(); + + const MaskMode mask_mode = static_cast(mask_mode_code); + + auto q_scalar_type = q_nope.scalar_type(); + auto kv_scalar_type = ckv_cache.scalar_type(); + + unsigned int q_nope_stride_n = q_nope.stride(0); + unsigned int q_nope_stride_h = q_nope.stride(1); + unsigned int q_pe_stride_n = q_pe.stride(0); + unsigned int q_pe_stride_h = q_pe.stride(1); + unsigned int ckv_stride_page = ckv_cache.stride(0); + unsigned int ckv_stride_n = ckv_cache.stride(1); + unsigned int kpe_stride_page = kpe_cache.stride(0); + unsigned int kpe_stride_n = kpe_cache.stride(1); + unsigned int o_stride_n = o.stride(0); + unsigned int o_stride_h = o.stride(1); + + const c10::cuda::OptionalCUDAGuard device_guard(device); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_CKV, HEAD_DIM_KPE, Params, [&] { + Params params; + + params.q_nope = static_cast(q_nope.data_ptr()); + params.q_pe = static_cast(q_pe.data_ptr()); + params.ckv = static_cast(ckv_cache.data_ptr()); + params.kpe = static_cast(kpe_cache.data_ptr()); + + params.q_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.q_indptr_offset); + params.kv_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_indptr_offset); + params.partial_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.partial_indptr_offset); + params.kv_indices = static_cast(kv_indices.data_ptr()); + params.q_len = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.q_len_offset); + params.kv_len = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_len_offset); + params.q_start = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.q_start_offset); + params.kv_start = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_start_offset); + params.kv_end = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_end_offset); + params.work_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.work_indptr_offset); + params.merge_packed_offset_start = GetPtrFromBaseOffset( + int_buffer_ptr, plan_info.merge_packed_offset_start_offset); + params.merge_packed_offset_end = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.merge_packed_offset_end_offset); + params.merge_partial_packed_offset_start = GetPtrFromBaseOffset( + int_buffer_ptr, plan_info.merge_partial_packed_offset_start_offset); + params.merge_partial_packed_offset_end = GetPtrFromBaseOffset( + int_buffer_ptr, plan_info.merge_partial_packed_offset_end_offset); + params.merge_partial_stride = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.merge_partial_stride_offset); + params.final_o = static_cast(o.data_ptr()); + params.final_lse = + maybe_lse.has_value() ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.partial_o = + GetPtrFromBaseOffset(float_buffer_ptr, plan_info.partial_o_offset); + params.partial_lse = + GetPtrFromBaseOffset(float_buffer_ptr, plan_info.partial_lse_offset); + + params.num_heads = uint_fastdiv(num_heads); + params.block_size = uint_fastdiv(page_size); + + params.q_nope_stride_n = q_nope_stride_n; + params.q_nope_stride_h = q_nope_stride_h; + params.q_pe_stride_n = q_pe_stride_n; + params.q_pe_stride_h = q_pe_stride_h; + params.ckv_stride_page = ckv_stride_page; + params.ckv_stride_n = ckv_stride_n; + params.kpe_stride_page = kpe_stride_page; + params.kpe_stride_n = kpe_stride_n; + params.o_stride_n = o_stride_n; + params.o_stride_h = o_stride_h; + + params.sm_scale = sm_scale; + + cudaError_t status = mla::BatchMLAPagedAttention( + params, plan_info.num_blks_x, plan_info.num_blks_y, stream); + + TORCH_CHECK(status == cudaSuccess, + "Failed to run MLA, error: ", cudaGetErrorString(status)); + }); +} diff --git a/csrc/generated/batch_mla_attention_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_ckv_512_head_dim_kpe_64_profiler_False_sm90/batch_mla_sm90_config.inc b/csrc/generated/batch_mla_attention_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_ckv_512_head_dim_kpe_64_profiler_False_sm90/batch_mla_sm90_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..6f4c28dc8ea7f7ab624ed810abb6ee6fe7cfccb7 --- /dev/null +++ b/csrc/generated/batch_mla_attention_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_ckv_512_head_dim_kpe_64_profiler_False_sm90/batch_mla_sm90_config.inc @@ -0,0 +1,33 @@ +#pragma once +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace flashinfer; + +#ifdef FLASHINFER_ENABLE_PROFILER +#define ADDITIONAL_FUNC_PARAMS , at::Tensor profiler_buffer +#define ADDITIONAL_PARAMS_SETTER \ + params.profiler_buffer = static_cast(profiler_buffer.data_ptr()); +#else +#define ADDITIONAL_FUNC_PARAMS +#define ADDITIONAL_PARAMS_SETTER +#endif + +using DTypeQ = half; +using DTypeKV = half; +using DTypeO = half; +using IdType = int32_t; +constexpr int HEAD_DIM_CKV = 512; +constexpr int HEAD_DIM_KPE = 64; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_CKV, HEAD_DIM_KPE, Params, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { \ + using Params = MLAParams; \ + __VA_ARGS__(); \ + }) \ No newline at end of file diff --git a/csrc/generated/batch_mla_attention_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_ckv_512_head_dim_kpe_64_profiler_False_sm90/batch_mla_sm90_plan.cu b/csrc/generated/batch_mla_attention_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_ckv_512_head_dim_kpe_64_profiler_False_sm90/batch_mla_sm90_plan.cu new file mode 100644 index 0000000000000000000000000000000000000000..35f6d0fa820ba335fdb04324c16f054ab66d2279 --- /dev/null +++ b/csrc/generated/batch_mla_attention_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_ckv_512_head_dim_kpe_64_profiler_False_sm90/batch_mla_sm90_plan.cu @@ -0,0 +1,53 @@ +/* + * Copyright (c) 2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include + +#include "batch_mla_sm90_config.inc" +#include "pytorch_conversion_utils.h" +#include "pytorch_extension_utils.h" + +using namespace flashinfer; + +at::Tensor BatchMLAPagedAttentionSM90Plan(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, + at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len, int64_t num_heads, int64_t head_dim_o, + bool causal) { + size_t float_workspace_size_in_bytes = + float_workspace_buffer.size(0) * float_workspace_buffer.element_size(); + size_t int_workspace_size_in_bytes = + int_workspace_buffer.size(0) * int_workspace_buffer.element_size(); + + MLAPlanInfo plan_info; + + int batch_size = kv_len.size(0); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + cudaError_t status = + MLAPlan(float_workspace_buffer.data_ptr(), float_workspace_size_in_bytes, + int_workspace_buffer.data_ptr(), page_locked_int_workspace_buffer.data_ptr(), + int_workspace_size_in_bytes, plan_info, static_cast(qo_indptr.data_ptr()), + static_cast(kv_indptr.data_ptr()), static_cast(kv_len.data_ptr()), + batch_size, num_heads, head_dim_o, causal, stream); + + TORCH_CHECK(status == cudaSuccess, "Failed to plan MLA, error: ", cudaGetErrorString(status)); + + return vec_to_tensor(plan_info.ToVector()); +} diff --git a/csrc/generated/batch_mla_attention_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_ckv_512_head_dim_kpe_64_profiler_False_sm90/batch_mla_sm90_pybind.cu b/csrc/generated/batch_mla_attention_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_ckv_512_head_dim_kpe_64_profiler_False_sm90/batch_mla_sm90_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..d089ad0c5c1acbd1c3ed788b8412d1ae0523ca49 --- /dev/null +++ b/csrc/generated/batch_mla_attention_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_ckv_512_head_dim_kpe_64_profiler_False_sm90/batch_mla_sm90_pybind.cu @@ -0,0 +1,37 @@ +/* + * Copyright (c) 2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "batch_mla_sm90_config.inc" +#include "pytorch_extension_utils.h" + +at::Tensor BatchMLAPagedAttentionSM90Plan(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, + at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len, int64_t num_heads, int64_t head_dim_o, + bool causal); + +void BatchMLAPagedAttentionSM90Run(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q_nope, at::Tensor q_pe, at::Tensor ckv_cache, + at::Tensor kpe_cache, at::Tensor kv_indices, at::Tensor o, + std::optional maybe_lse, int64_t mask_mode_code, + int64_t num_heads, int64_t page_size, + double sm_scale ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + m.def("plan", &BatchMLAPagedAttentionSM90Plan); + m.def("run", &BatchMLAPagedAttentionSM90Run); +} diff --git a/csrc/generated/batch_mla_attention_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_ckv_512_head_dim_kpe_64_profiler_False_sm90/batch_mla_sm90_run.cu b/csrc/generated/batch_mla_attention_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_ckv_512_head_dim_kpe_64_profiler_False_sm90/batch_mla_sm90_run.cu new file mode 100644 index 0000000000000000000000000000000000000000..3da072eb039d16a3cf7ee319be03c487ac6fc288 --- /dev/null +++ b/csrc/generated/batch_mla_attention_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_ckv_512_head_dim_kpe_64_profiler_False_sm90/batch_mla_sm90_run.cu @@ -0,0 +1,129 @@ +/* + * Copyright (c) 2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include +#include + +#include "batch_mla_sm90_config.inc" +#include "pytorch_conversion_utils.h" +#include "pytorch_extension_utils.h" + +using namespace flashinfer; + +void BatchMLAPagedAttentionSM90Run(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q_nope, at::Tensor q_pe, at::Tensor ckv_cache, + at::Tensor kpe_cache, at::Tensor kv_indices, at::Tensor o, + std::optional maybe_lse, int64_t mask_mode_code, + int64_t num_heads, int64_t page_size, + double sm_scale ADDITIONAL_FUNC_PARAMS) { + // q_nope: [n, num_heads, head_dim_ckv] + // q_pe: [n, num_heads, head_dim_kpe] + // ckv_cache: [num_pages, page_size, head_dim_ckv] + // kpe_cache: [num_pages, page_size, head_dim_kpe] + MLAPlanInfo plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + + auto device = q_nope.device(); + + void* float_buffer_ptr = float_workspace_buffer.data_ptr(); + void* int_buffer_ptr = int_workspace_buffer.data_ptr(); + + const MaskMode mask_mode = static_cast(mask_mode_code); + + auto q_scalar_type = q_nope.scalar_type(); + auto kv_scalar_type = ckv_cache.scalar_type(); + + unsigned int q_nope_stride_n = q_nope.stride(0); + unsigned int q_nope_stride_h = q_nope.stride(1); + unsigned int q_pe_stride_n = q_pe.stride(0); + unsigned int q_pe_stride_h = q_pe.stride(1); + unsigned int ckv_stride_page = ckv_cache.stride(0); + unsigned int ckv_stride_n = ckv_cache.stride(1); + unsigned int kpe_stride_page = kpe_cache.stride(0); + unsigned int kpe_stride_n = kpe_cache.stride(1); + unsigned int o_stride_n = o.stride(0); + unsigned int o_stride_h = o.stride(1); + + const c10::cuda::OptionalCUDAGuard device_guard(device); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_CKV, HEAD_DIM_KPE, Params, [&] { + Params params; + + params.q_nope = static_cast(q_nope.data_ptr()); + params.q_pe = static_cast(q_pe.data_ptr()); + params.ckv = static_cast(ckv_cache.data_ptr()); + params.kpe = static_cast(kpe_cache.data_ptr()); + + params.q_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.q_indptr_offset); + params.kv_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_indptr_offset); + params.partial_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.partial_indptr_offset); + params.kv_indices = static_cast(kv_indices.data_ptr()); + params.q_len = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.q_len_offset); + params.kv_len = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_len_offset); + params.q_start = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.q_start_offset); + params.kv_start = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_start_offset); + params.kv_end = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_end_offset); + params.work_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.work_indptr_offset); + params.merge_packed_offset_start = GetPtrFromBaseOffset( + int_buffer_ptr, plan_info.merge_packed_offset_start_offset); + params.merge_packed_offset_end = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.merge_packed_offset_end_offset); + params.merge_partial_packed_offset_start = GetPtrFromBaseOffset( + int_buffer_ptr, plan_info.merge_partial_packed_offset_start_offset); + params.merge_partial_packed_offset_end = GetPtrFromBaseOffset( + int_buffer_ptr, plan_info.merge_partial_packed_offset_end_offset); + params.merge_partial_stride = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.merge_partial_stride_offset); + params.final_o = static_cast(o.data_ptr()); + params.final_lse = + maybe_lse.has_value() ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.partial_o = + GetPtrFromBaseOffset(float_buffer_ptr, plan_info.partial_o_offset); + params.partial_lse = + GetPtrFromBaseOffset(float_buffer_ptr, plan_info.partial_lse_offset); + + params.num_heads = uint_fastdiv(num_heads); + params.block_size = uint_fastdiv(page_size); + + params.q_nope_stride_n = q_nope_stride_n; + params.q_nope_stride_h = q_nope_stride_h; + params.q_pe_stride_n = q_pe_stride_n; + params.q_pe_stride_h = q_pe_stride_h; + params.ckv_stride_page = ckv_stride_page; + params.ckv_stride_n = ckv_stride_n; + params.kpe_stride_page = kpe_stride_page; + params.kpe_stride_n = kpe_stride_n; + params.o_stride_n = o_stride_n; + params.o_stride_h = o_stride_h; + + ADDITIONAL_PARAMS_SETTER + + params.sm_scale = sm_scale; + + cudaError_t status = + mla::BatchMLAPageAttentionHopper( + params, plan_info.num_blks_x, plan_info.num_blks_y, stream); + + TORCH_CHECK(status == cudaSuccess, + "Failed to run MLA, error: ", cudaGetErrorString(status)); + }); +} diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill.cu new file mode 100644 index 0000000000000000000000000000000000000000..a51fc7f56a0af7482bd5d2a9d3f672854df69dfd --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill.cu @@ -0,0 +1,330 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include +#include + +#include "batch_prefill_config.inc" +#include "pytorch_conversion_utils.h" +#include "pytorch_extension_utils.h" + +namespace flashinfer { + +template +cudaError_t BatchPrefillWithPagedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v, + float* tmp_s, bool enable_pdl, + cudaStream_t stream); + +template +cudaError_t BatchPrefillWithRaggedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v, + float* tmp_s, bool enable_pdl, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +at::Tensor BatchPrefillWithKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, + int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, + int64_t head_dim_vo, bool causal) { + size_t float_workspace_size_in_bytes = + float_workspace_buffer.size(0) * float_workspace_buffer.element_size(); + size_t int_workspace_size_in_bytes = + int_workspace_buffer.size(0) * int_workspace_buffer.element_size(); + + PrefillPlanInfo plan_info; + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + cudaError_t status = PrefillPlan( + float_workspace_buffer.data_ptr(), float_workspace_size_in_bytes, + int_workspace_buffer.data_ptr(), page_locked_int_workspace_buffer.data_ptr(), + int_workspace_size_in_bytes, plan_info, qo_indptr.data_ptr(), + kv_indptr.data_ptr(), total_num_rows, batch_size, num_qo_heads, num_kv_heads, + head_dim_qk, head_dim_vo, page_size, enable_cuda_graph, /*sizeof_dtype_o=*/2, stream); + + TORCH_CHECK(status == cudaSuccess, + "Failed to plan prefill with error: ", cudaGetErrorString(status)); + + return vec_to_tensor(plan_info.ToVector()); +} + +void BatchPrefillWithRaggedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor k, at::Tensor v, + at::Tensor qo_indptr, at::Tensor kv_indptr, at::Tensor o, + std::optional maybe_lse, int64_t mask_mode_code, + int64_t layout, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + PrefillPlanInfo plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + QKVLayout kv_layout = static_cast(layout); + + int64_t num_qo_heads = q.size(1); + int64_t head_dim_qk = q.size(2); + int64_t num_kv_heads = (kv_layout == QKVLayout::kNHD) ? k.size(1) : k.size(0); + uint32_t q_stride_n = q.stride(0), q_stride_h = q.stride(1), k_stride_n, k_stride_h, v_stride_n, + v_stride_h; + if (kv_layout == QKVLayout::kNHD) { + k_stride_n = k.stride(0); + k_stride_h = k.stride(1); + v_stride_n = v.stride(0); + v_stride_h = v.stride(1); + } else { + k_stride_h = k.stride(0); + k_stride_n = k.stride(1); + v_stride_h = v.stride(0); + v_stride_n = v.stride(1); + } + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == q.size(0), lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == q.size(1), lse.size(1), q.size(1)); + } + + void* float_buffer_ptr = float_workspace_buffer.data_ptr(); + void* int_buffer_ptr = int_workspace_buffer.data_ptr(); + + const MaskMode mask_mode = static_cast(mask_mode_code); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, + RaggedParams, PagedParams, [&] { + RaggedParams params; + + params.q = static_cast(q.data_ptr()); + params.k = static_cast(k.data_ptr()); + params.v = static_cast(v.data_ptr()); + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.q_indptr = static_cast(qo_indptr.data_ptr()); + params.kv_indptr = static_cast(kv_indptr.data_ptr()); + params.num_qo_heads = num_qo_heads; + params.num_kv_heads = num_kv_heads; + params.group_size = uint_fastdiv(num_qo_heads / num_kv_heads); + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.k_stride_n = k_stride_n; + params.k_stride_h = k_stride_h; + params.v_stride_n = v_stride_n; + params.v_stride_h = v_stride_h; + params.window_left = window_left; + + params.request_indices = nullptr; + params.qo_tile_indices = nullptr; + params.kv_tile_indices = nullptr; + params.merge_indptr = nullptr; + params.o_indptr = nullptr; + params.kv_chunk_size_ptr = nullptr; + params.block_valid_mask = nullptr; + params.total_num_rows = nullptr; + params.max_total_num_rows = 0; + params.padded_batch_size = 0; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + DTypeO* tmp_v = nullptr; + float* tmp_s = nullptr; + + params.request_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.request_indices_offset); + params.qo_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_tile_indices_offset); + params.kv_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_tile_indices_offset); + params.o_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.o_indptr_offset); + params.kv_chunk_size_ptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_chunk_size_ptr_offset); + if (plan_info.split_kv) { + params.merge_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.merge_indptr_offset); + tmp_v = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.v_offset); + tmp_s = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.s_offset); + if (plan_info.enable_cuda_graph) { + params.block_valid_mask = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.block_valid_mask_offset); + } + } + params.padded_batch_size = plan_info.padded_batch_size; + params.max_total_num_rows = plan_info.total_num_rows; + if (plan_info.enable_cuda_graph) { + params.total_num_rows = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.total_num_rows_offset); + } + + cudaError_t status = cudaSuccess; + + DISPATCH_CTA_TILE_Q(plan_info.cta_tile_q, CTA_TILE_Q, { + status = flashinfer::BatchPrefillWithRaggedKVCacheDispatched< + CTA_TILE_Q, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + /*use_fp16_qk_reduction=*/USE_FP16_QK_REDUCTION, MASK_MODE, AttentionVariant, + RaggedParams>(params, tmp_v, tmp_s, enable_pdl, stream); + }); + + TORCH_CHECK(status == cudaSuccess, "BatchPrefillWithRaggedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); +} + +void BatchPrefillWithPagedKVCacheRun( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, + at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + PrefillPlanInfo plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + QKVLayout kv_layout = static_cast(layout); + auto device = q.device(); + int64_t batch_size = paged_kv_indptr.size(0) - 1; + int64_t num_qo_heads = q.size(1); + int64_t num_kv_heads, page_size; + uint32_t head_dim_qk = q.size(2); + if (kv_layout == QKVLayout::kHND) { + num_kv_heads = paged_k_cache.size(1); + page_size = paged_k_cache.size(2); + } else { + page_size = paged_k_cache.size(1); + num_kv_heads = paged_k_cache.size(2); + } + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == q.size(0), lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == q.size(1), lse.size(1), q.size(1)); + } + + void* float_buffer_ptr = static_cast(float_workspace_buffer.data_ptr()); + void* int_buffer_ptr = static_cast(int_workspace_buffer.data_ptr()); + + const MaskMode mask_mode = static_cast(mask_mode_code); + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = paged_k_cache.scalar_type(); + + // get q_stride_n and q_stride_h + const auto q_stride_n = q.stride(0); + const auto q_stride_h = q.stride(1); + + // get kv_cache_strides + const int64_t* kv_cache_strides = nullptr; + auto k_strides = paged_k_cache.strides(); + auto v_strides = paged_v_cache.strides(); + TORCH_CHECK(k_strides == v_strides, "k/v strides must be identical"); + kv_cache_strides = k_strides.data(); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, + RaggedParams, PagedParams, [&] { + PagedParams params; + + params.q = static_cast(q.data_ptr()); + paged_kv_t paged_kv( + num_kv_heads, page_size, HEAD_DIM_VO, batch_size, kv_layout, + static_cast(paged_k_cache.data_ptr()), + static_cast(paged_v_cache.data_ptr()), kv_cache_strides, + static_cast(paged_kv_indices.data_ptr()), + static_cast(paged_kv_indptr.data_ptr()), + static_cast(paged_kv_last_page_len.data_ptr())); + params.paged_kv = paged_kv; + params.q_indptr = static_cast(qo_indptr.data_ptr()); + params.o = static_cast(o.data_ptr()); + + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.num_qo_heads = num_qo_heads; + params.group_size = uint_fastdiv(num_qo_heads / paged_kv.num_heads); + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.window_left = window_left; + + params.request_indices = nullptr; + params.qo_tile_indices = nullptr; + params.kv_tile_indices = nullptr; + params.merge_indptr = nullptr; + params.o_indptr = nullptr; + params.kv_chunk_size_ptr = nullptr; + params.block_valid_mask = nullptr; + params.total_num_rows = nullptr; + params.max_total_num_rows = 0; + params.padded_batch_size = 0; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + DTypeO* tmp_v = nullptr; + float* tmp_s = nullptr; + + params.request_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.request_indices_offset); + params.qo_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_tile_indices_offset); + params.kv_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_tile_indices_offset); + params.o_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.o_indptr_offset); + params.kv_chunk_size_ptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_chunk_size_ptr_offset); + if (plan_info.split_kv) { + params.merge_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.merge_indptr_offset); + tmp_v = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.v_offset); + tmp_s = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.s_offset); + if (plan_info.enable_cuda_graph) { + params.block_valid_mask = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.block_valid_mask_offset); + } + } + params.padded_batch_size = plan_info.padded_batch_size; + params.max_total_num_rows = plan_info.total_num_rows; + if (plan_info.enable_cuda_graph) { + params.total_num_rows = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.total_num_rows_offset); + } + + cudaError_t status = cudaSuccess; + + DISPATCH_CTA_TILE_Q(plan_info.cta_tile_q, CTA_TILE_Q, { + status = flashinfer::BatchPrefillWithPagedKVCacheDispatched< + CTA_TILE_Q, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + /*use_fp16_qk_reduction=*/USE_FP16_QK_REDUCTION, MASK_MODE, AttentionVariant, + PagedParams>(params, tmp_v, tmp_s, enable_pdl, stream); + }); + + TORCH_CHECK(status == cudaSuccess, "BatchPrefillWithPagedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); +} diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_config.inc b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..caa0a075a5e124f7c516a4fe43ff6588f9c7c675 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_config.inc @@ -0,0 +1,153 @@ +#pragma once +#include +#include +#include +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_custom_mask, std::optional maybe_mask_indptr, std::optional maybe_alibi_slopes, std::optional maybe_prefix_len_ptr, std::optional maybe_token_pos_in_items_ptr, std::optional maybe_max_item_len_ptr, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta, int64_t token_pos_in_items_len +#define ADDITIONAL_PARAMS_SETTER params.maybe_custom_mask = maybe_custom_mask ? static_cast(maybe_custom_mask->data_ptr()): nullptr; \ +params.maybe_mask_indptr = maybe_mask_indptr ? static_cast(maybe_mask_indptr->data_ptr()): nullptr; \ +params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.maybe_prefix_len_ptr = maybe_prefix_len_ptr ? static_cast(maybe_prefix_len_ptr->data_ptr()): nullptr; \ +params.maybe_token_pos_in_items_ptr = maybe_token_pos_in_items_ptr ? static_cast(maybe_token_pos_in_items_ptr->data_ptr()): nullptr; \ +params.maybe_max_item_len_ptr = maybe_max_item_len_ptr ? static_cast(maybe_max_item_len_ptr->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; \ +params.token_pos_in_items_len = token_pos_in_items_len; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, RaggedParams, PagedParams, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { \ + constexpr auto use_custom_mask = MASK_MODE == MaskMode::kCustom; \ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ + }) + +using namespace flashinfer; + +using DTypeQ = nv_bfloat16; +using DTypeKV = nv_bfloat16; +using DTypeO = nv_bfloat16; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 128; +constexpr int HEAD_DIM_VO = 128; +constexpr bool USE_FP16_QK_REDUCTION = false; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = false; + + +struct RaggedParams { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + DTypeQ* q; + DTypeKV* k; + DTypeKV* v; + IdType* q_indptr; + IdType* kv_indptr; + DTypeO* o; + float* lse; + uint_fastdiv group_size; + + uint8_t* maybe_custom_mask; +int32_t* maybe_mask_indptr; +float* maybe_alibi_slopes; +uint32_t* maybe_prefix_len_ptr; +uint16_t* maybe_token_pos_in_items_ptr; +uint16_t* maybe_max_item_len_ptr; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; +int64_t token_pos_in_items_len; + + uint32_t num_qo_heads; + uint32_t num_kv_heads; + uint32_t q_stride_n; + uint32_t q_stride_h; + uint32_t k_stride_n; + uint32_t k_stride_h; + uint32_t v_stride_n; + uint32_t v_stride_h; + int32_t window_left; + + IdType* request_indices; + IdType* qo_tile_indices; + IdType* kv_tile_indices; + IdType* merge_indptr; + IdType* o_indptr; + IdType* kv_chunk_size_ptr; + bool* block_valid_mask; + uint32_t max_total_num_rows; + uint32_t* total_num_rows; + uint32_t padded_batch_size; + bool partition_kv; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { + return q_indptr[batch_idx + 1] - q_indptr[batch_idx]; + } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return kv_indptr[batch_idx + 1] - kv_indptr[batch_idx]; + } +}; + +struct PagedParams { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + DTypeQ* q; + paged_kv_t paged_kv; + IdType* q_indptr; + DTypeO* o; + float* lse; + uint_fastdiv group_size; + + uint8_t* maybe_custom_mask; +int32_t* maybe_mask_indptr; +float* maybe_alibi_slopes; +uint32_t* maybe_prefix_len_ptr; +uint16_t* maybe_token_pos_in_items_ptr; +uint16_t* maybe_max_item_len_ptr; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; +int64_t token_pos_in_items_len; + + uint32_t num_qo_heads; + IdType q_stride_n; + IdType q_stride_h; + int32_t window_left; + + IdType* request_indices; + IdType* qo_tile_indices; + IdType* kv_tile_indices; + IdType* merge_indptr; + IdType* o_indptr; + bool* block_valid_mask; + IdType* kv_chunk_size_ptr; + uint32_t max_total_num_rows; + uint32_t* total_num_rows; + uint32_t padded_batch_size; + bool partition_kv; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { + return q_indptr[batch_idx + 1] - q_indptr[batch_idx]; + } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return paged_kv.get_length(batch_idx); + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_jit_pybind.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..5421ab1cf910fdfbc867c182c8e22f23e44749ae --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_jit_pybind.cu @@ -0,0 +1,48 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "batch_prefill_config.inc" +#include "pytorch_extension_utils.h" + +at::Tensor BatchPrefillWithKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, + int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, + int64_t head_dim_vo, bool causal); + +void BatchPrefillWithRaggedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor k, at::Tensor v, + at::Tensor qo_indptr, at::Tensor kv_indptr, at::Tensor o, + std::optional maybe_lse, int64_t mask_mode_code, + int64_t layout, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +void BatchPrefillWithPagedKVCacheRun( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, + at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Batch-request prefill attention with KV-Cache plan + m.def("plan", BatchPrefillWithKVCachePlan); + // Batch-request prefill attention with KV-Cache operator + m.def("ragged_run", BatchPrefillWithRaggedKVCacheRun); + // Batch-request prefill attention with KV-Cache operator + m.def("paged_run", BatchPrefillWithPagedKVCacheRun); +} diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_0.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..4c306789cc8cbc7b47b4533007bd6ad51ab9a9e1 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_0.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kNone == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 128, 128, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 128, 128, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 128, 128, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_1.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..d40fe782b6dd930e6e8548d9c15df1dacaee3b8b --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_1.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCausal == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 128, 128, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 128, 128, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 128, 128, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_2.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..16e8d67b528952a4650148a9e7119057667f18f6 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_2.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCustom == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 128, 128, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 128, 128, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 128, 128, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_3.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..7faebee71230224faa5fba4ec9709f8d5f811ca2 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_3.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kMultiItemScoring == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 128, 128, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 128, 128, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 128, 128, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_0.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..38d25465ee0a5e231720942d0712969244567fbe --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_0.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kNone == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 128, 128, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 128, 128, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 128, 128, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_1.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..f18906355ac04ee88e405d998e5d3aa211884b1c --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_1.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCausal == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 128, 128, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 128, 128, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 128, 128, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_2.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..bf34ea033aefab899927a75acc5ed1a6e7a7b310 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_2.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCustom == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 128, 128, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 128, 128, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 128, 128, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_3.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..6b6e36ad0a7464e476520a6e6761370b217dfcbb --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_3.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kMultiItemScoring == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 128, 128, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 128, 128, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 128, 128, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_0.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..7a3dcb47ab89daa3289551969a12c2a2c53c3d19 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_0.cu @@ -0,0 +1,24 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched + <128, + 128, + MaskMode::kNone, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultAttention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched + <128, + 128, + MaskMode::kNone, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultAttention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_1.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..7374938760f2ca572a9bcee91465b15daf194e5d --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_1.cu @@ -0,0 +1,24 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched + <128, + 128, + MaskMode::kCausal, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultAttention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched + <128, + 128, + MaskMode::kCausal, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultAttention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_2.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..522184fb2d414cba8ce56e576d475fee0651f912 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_2.cu @@ -0,0 +1,24 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched + <128, + 128, + MaskMode::kCustom, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultAttention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched + <128, + 128, + MaskMode::kCustom, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultAttention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_3.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..4a657661a8d94cbdc5c4ee19bb75e13db7588a2b --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_3.cu @@ -0,0 +1,24 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched + <128, + 128, + MaskMode::kMultiItemScoring, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultAttention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched + <128, + 128, + MaskMode::kMultiItemScoring, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultAttention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_0.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..bf033b46b81da64fc48ade098472326bd5c2712f --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_0.cu @@ -0,0 +1,24 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched + <128, + 128, + MaskMode::kNone, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultAttention>(RaggedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched + <128, + 128, + MaskMode::kNone, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultAttention>(RaggedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_1.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..51e21ba99d82e9e4d3840ea6dc1ee260d7550580 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_1.cu @@ -0,0 +1,24 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched + <128, + 128, + MaskMode::kCausal, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultAttention>(RaggedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched + <128, + 128, + MaskMode::kCausal, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultAttention>(RaggedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_2.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..0c7e6867e3aed1752226dc283c266ed97682e290 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_2.cu @@ -0,0 +1,24 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched + <128, + 128, + MaskMode::kCustom, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultAttention>(RaggedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched + <128, + 128, + MaskMode::kCustom, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultAttention>(RaggedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_3.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..e15debd96911a4e77e702830b7c4d414410e3e98 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_3.cu @@ -0,0 +1,24 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched + <128, + 128, + MaskMode::kMultiItemScoring, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultAttention>(RaggedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched + <128, + 128, + MaskMode::kMultiItemScoring, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultAttention>(RaggedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_sm90.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_sm90.cu new file mode 100644 index 0000000000000000000000000000000000000000..159ee3f1eda3f47769ec14f8fac16f35fa9bda05 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_sm90.cu @@ -0,0 +1,260 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include +#include + +#include "batch_prefill_sm90_config.inc" +#include "pytorch_conversion_utils.h" +#include "pytorch_extension_utils.h" + +namespace flashinfer { + +template +cudaError_t BatchPrefillWithRaggedKVCacheDispatched(Params& params, bool enable_pdl, + cudaStream_t stream); + +template +cudaError_t BatchPrefillWithPagedKVCacheDispatched(Params& params, bool enable_pdl, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +at::Tensor BatchPrefillWithKVCacheSM90Plan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, + int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, + int64_t head_dim_vo, bool causal) { + size_t float_workspace_size_in_bytes = + float_workspace_buffer.size(0) * float_workspace_buffer.element_size(); + size_t int_workspace_size_in_bytes = + int_workspace_buffer.size(0) * int_workspace_buffer.element_size(); + + flashinfer::PrefillPlanSM90Info plan_info; + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + cudaError_t status = + PrefillSM90Plan(float_workspace_buffer.data_ptr(), float_workspace_size_in_bytes, + int_workspace_buffer.data_ptr(), page_locked_int_workspace_buffer.data_ptr(), + int_workspace_size_in_bytes, plan_info, qo_indptr.data_ptr(), + kv_indptr.data_ptr(), kv_len_arr.data_ptr(), total_num_rows, + batch_size, num_qo_heads, num_kv_heads, head_dim_qk, head_dim_vo, page_size, + causal, enable_cuda_graph, /*sizeof_dtype_o=*/2, stream); + + TORCH_CHECK(status == cudaSuccess, + "PrefillSM90Plan failed with error: ", cudaGetErrorString(status)); + + return vec_to_tensor(plan_info.ToVector()); +} + +void BatchPrefillWithRaggedKVCacheSM90Run( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + PrefillPlanSM90Info plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == q.size(0), lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == q.size(1), lse.size(1), q.size(1)); + } + + void* float_buffer_ptr = float_workspace_buffer.data_ptr(); + void* int_buffer_ptr = int_workspace_buffer.data_ptr(); + + int64_t head_dim_qk = q.size(2); + int64_t head_dim_vo = v.size(2); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + QKVLayout kv_layout = static_cast(layout); + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + const MaskMode mask_mode = static_cast(mask_mode_code); + bool use_swa = window_left != -1; + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, USE_SLIDING_WINDOW, + USE_LOGITS_SOFT_CAP, AttentionVariant, RaggedParams, PagedParams, [&] { + RaggedParams params; + + params.q_ptr = static_cast(q.data_ptr()); + params.k_ptr = static_cast(k.data_ptr()); + params.v_ptr = static_cast(v.data_ptr()); + params.o_ptr = static_cast(o.data_ptr()); + params.lse_ptr = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.q_stride_n = q.stride(0); + params.q_stride_h = q.stride(1); + params.o_stride_n = o.stride(0); + params.o_stride_h = o.stride(1); + if (kv_layout == QKVLayout::kNHD) { + params.k_stride_n = k.stride(0); + params.k_stride_h = k.stride(1); + params.v_stride_n = v.stride(0); + params.v_stride_h = v.stride(1); + } else { + params.k_stride_h = k.stride(0); + params.k_stride_n = k.stride(1); + params.v_stride_h = v.stride(0); + params.v_stride_n = v.stride(1); + } + params.nnz_qo = q.size(0); + params.nnz_kv = k.size(0); + params.num_qo_heads = q.size(1); + params.num_kv_heads = k.size(1); + params.group_size = params.num_qo_heads / params.num_kv_heads; + params.window_left = window_left; + params.causal = mask_mode_code == 1; + params.qo_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_tile_indices_offset); + params.qo_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_indptr_offset); + params.kv_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_indptr_offset); + params.qo_lens = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_len_offset); + params.kv_lens = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_len_offset); + params.head_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.head_indices_offset); + params.work_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.work_indptr_offset); + params.batch_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.batch_indices_offset); + + ADDITIONAL_PARAMS_SETTER + + bool same_schedule_for_all_heads = plan_info.same_schedule_for_all_heads; + DISPATCH_BOOL(same_schedule_for_all_heads, SAME_SCHEDULER_FOR_ALL_HEADS, [&] { + cudaError_t status = BatchPrefillWithRaggedKVCacheDispatched< + HEAD_DIM_QK, HEAD_DIM_VO, MASK_MODE, USE_SLIDING_WINDOW, SAME_SCHEDULER_FOR_ALL_HEADS, + AttentionVariant>(params, enable_pdl, stream); + TORCH_CHECK(status == cudaSuccess, + "BatchPrefillWithRaggedKVCacheSM90Run failed with error: ", + cudaGetErrorString(status)); + return true; + }); + }); +} + +void BatchPrefillWithPagedKVCacheSM90Run( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, + at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + PrefillPlanSM90Info plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == q.size(0), lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == q.size(1), lse.size(1), q.size(1)); + } + QKVLayout kv_layout = static_cast(layout); + int64_t num_kv_heads, page_size; + int64_t head_dim_qk = q.size(2); + int64_t head_dim_vo = paged_v_cache.size(3); + if (kv_layout == QKVLayout::kHND) { + num_kv_heads = paged_k_cache.size(1); + page_size = paged_k_cache.size(2); + } else { + page_size = paged_k_cache.size(1); + num_kv_heads = paged_k_cache.size(2); + } + + void* float_buffer_ptr = float_workspace_buffer.data_ptr(); + void* int_buffer_ptr = int_workspace_buffer.data_ptr(); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = paged_k_cache.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + const MaskMode mask_mode = static_cast(mask_mode_code); + bool use_swa = window_left != -1; + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, USE_SLIDING_WINDOW, + USE_LOGITS_SOFT_CAP, AttentionVariant, RaggedParams, PagedParams, [&] { + PagedParams params; + + params.q_ptr = static_cast(q.data_ptr()); + params.k_ptr = static_cast(paged_k_cache.data_ptr()); + params.v_ptr = static_cast(paged_v_cache.data_ptr()); + params.o_ptr = static_cast(o.data_ptr()); + params.lse_ptr = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.q_stride_n = q.stride(0); + params.q_stride_h = q.stride(1); + params.o_stride_n = o.stride(0); + params.o_stride_h = o.stride(1); + if (kv_layout == QKVLayout::kNHD) { + // (num_pages, page_size, num_heads, head_dim) + params.k_stride_n = paged_k_cache.stride(1); + params.k_stride_h = paged_k_cache.stride(2); + params.v_stride_n = paged_v_cache.stride(1); + params.v_stride_h = paged_v_cache.stride(2); + } else { + // (num_pages, num_heads, page_size, head_dim) + params.k_stride_h = paged_k_cache.stride(1); + params.k_stride_n = paged_k_cache.stride(2); + params.v_stride_h = paged_v_cache.stride(1); + params.v_stride_n = paged_v_cache.stride(2); + } + params.nnz_qo = q.size(0); + params.num_qo_heads = q.size(1); + params.num_kv_heads = num_kv_heads; + params.group_size = params.num_qo_heads / num_kv_heads; + params.page_size = page_size; + params.window_left = window_left; + params.causal = mask_mode_code == 1; + params.qo_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_tile_indices_offset); + params.qo_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_indptr_offset); + params.kv_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_indptr_offset); + params.qo_lens = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_len_offset); + params.kv_lens = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_len_offset); + params.head_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.head_indices_offset); + params.work_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.work_indptr_offset); + params.batch_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.batch_indices_offset); + params.kv_indices = static_cast(paged_kv_indices.data_ptr()); + + ADDITIONAL_PARAMS_SETTER + + bool same_schedule_for_all_heads = plan_info.same_schedule_for_all_heads; + DISPATCH_BOOL(same_schedule_for_all_heads, SAME_SCHEDULER_FOR_ALL_HEADS, [&] { + cudaError_t status = BatchPrefillWithPagedKVCacheDispatched< + HEAD_DIM_QK, HEAD_DIM_VO, MASK_MODE, USE_SLIDING_WINDOW, SAME_SCHEDULER_FOR_ALL_HEADS, + AttentionVariant>(params, enable_pdl, stream); + TORCH_CHECK(status == cudaSuccess, + "BatchPrefillWithPagedKVCacheSM90Run failed with error: ", + cudaGetErrorString(status)); + return true; + }); + }); +} diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_sm90_config.inc b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_sm90_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..cdb5127361567bd8b985cd894dba4023c74a4bff --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_sm90_config.inc @@ -0,0 +1,135 @@ +#pragma once +#include +#include +#include +#include +#include +#include "pytorch_extension_utils.h" + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_prefix_len_ptr, std::optional maybe_token_pos_in_items_ptr, std::optional maybe_max_item_len_ptr, double logits_soft_cap, double sm_scale, int64_t token_pos_in_items_len +#define ADDITIONAL_PARAMS_SETTER params.additional_params.maybe_prefix_len_ptr = maybe_prefix_len_ptr ? static_cast(maybe_prefix_len_ptr->data_ptr()): nullptr; \ +params.additional_params.maybe_token_pos_in_items_ptr = maybe_token_pos_in_items_ptr ? static_cast(maybe_token_pos_in_items_ptr->data_ptr()): nullptr; \ +params.additional_params.maybe_max_item_len_ptr = maybe_max_item_len_ptr ? static_cast(maybe_max_item_len_ptr->data_ptr()): nullptr; \ +params.additional_params.logits_soft_cap = logits_soft_cap; \ +params.additional_params.sm_scale = sm_scale; \ +params.additional_params.token_pos_in_items_len = token_pos_in_items_len; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, RaggedParams, PagedParams, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { using AttentionVariant = DefaultAttention; __VA_ARGS__();}) + +using namespace flashinfer; + +using DTypeQ = cutlass_dtype_t; +using DTypeKV = cutlass_dtype_t; +using DTypeO = cutlass_dtype_t; +using IdType = cutlass_dtype_t; + +constexpr int HEAD_DIM_QK = 128; +constexpr int HEAD_DIM_VO = 128; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto USE_SLIDING_WINDOW = false; + +struct RaggedParams { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + // The QKV matrices. + DTypeQ* q_ptr; + DTypeKV* k_ptr; + DTypeKV* v_ptr; + DTypeO* o_ptr; + float* lse_ptr; + + IdType* qo_tile_indices; + IdType* qo_indptr; + IdType* kv_indptr; + IdType* qo_lens; + IdType* kv_lens; + IdType* head_indices; + IdType* work_indptr; + IdType* batch_indices; + + struct AdditionalParams { + uint32_t* maybe_prefix_len_ptr; +uint16_t* maybe_token_pos_in_items_ptr; +uint16_t* maybe_max_item_len_ptr; +double logits_soft_cap; +double sm_scale; +int64_t token_pos_in_items_len; + + } additional_params; + + int64_t q_stride_n; + int64_t k_stride_n; + int64_t v_stride_n; + int64_t o_stride_n; + int64_t q_stride_h; + int64_t k_stride_h; + int64_t v_stride_h; + int64_t o_stride_h; + int64_t nnz_qo; + int64_t nnz_kv; + + int head_dim; + int num_qo_heads; + int num_kv_heads; + int group_size; + int window_left; + + bool causal; +}; + +struct PagedParams { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + // The QKV matrices. + DTypeQ* q_ptr; + DTypeKV* k_ptr; + DTypeKV* v_ptr; + DTypeO* o_ptr; + float* lse_ptr; + + IdType* qo_tile_indices; + IdType* qo_indptr; + IdType* kv_indptr; + IdType* kv_indices; + IdType* qo_lens; + IdType* kv_lens; + IdType* head_indices; + IdType* work_indptr; + IdType* batch_indices; + + struct AdditionalParams { + uint32_t* maybe_prefix_len_ptr; +uint16_t* maybe_token_pos_in_items_ptr; +uint16_t* maybe_max_item_len_ptr; +double logits_soft_cap; +double sm_scale; +int64_t token_pos_in_items_len; + + } additional_params; + + int64_t q_stride_n; + int64_t k_stride_n; + int64_t v_stride_n; + int64_t o_stride_n; + int64_t q_stride_h; + int64_t k_stride_h; + int64_t v_stride_h; + int64_t o_stride_h; + int64_t nnz_qo; + + int head_dim; + int num_qo_heads; + int num_kv_heads; + int group_size; + int page_size; + int window_left; + + bool causal; +}; + +#include \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_sm90_jit_pybind.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_sm90_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..b1edd582549e124c776f2cf9b4503f4bf60d16ae --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_sm90_jit_pybind.cu @@ -0,0 +1,46 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "batch_prefill_sm90_config.inc" +#include "pytorch_extension_utils.h" + +at::Tensor BatchPrefillWithKVCacheSM90Plan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, + int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, + int64_t head_dim_vo, bool causal); + +void BatchPrefillWithRaggedKVCacheSM90Run( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +void BatchPrefillWithPagedKVCacheSM90Run( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, + at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Batch-request prefill attention with KV-Cache plan + m.def("plan", BatchPrefillWithKVCacheSM90Plan); + // Batch-request prefill attention with KV-Cache operator + m.def("ragged_run", BatchPrefillWithRaggedKVCacheSM90Run); + // Batch-request prefill attention with KV-Cache operator + m.def("paged_run", BatchPrefillWithPagedKVCacheSM90Run); +} diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_0.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..857a0f5796850a0d4499248335a5c7f1e2dad25b --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_0.cu @@ -0,0 +1,24 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched + <192, + 128, + MaskMode::kNone, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultAttention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched + <192, + 128, + MaskMode::kNone, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultAttention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_1.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..2f298fe5b33cea4a229c62692d707ea5f8d6fce5 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_1.cu @@ -0,0 +1,24 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched + <192, + 128, + MaskMode::kCausal, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultAttention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched + <192, + 128, + MaskMode::kCausal, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultAttention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_2.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..d075a56a378fce5b3b4c3a14e75309f3db9e560f --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_2.cu @@ -0,0 +1,24 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched + <192, + 128, + MaskMode::kCustom, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultAttention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched + <192, + 128, + MaskMode::kCustom, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultAttention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_3.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..8f188b661bb83a6220f002a64d7f2bd0d91d482f --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_3.cu @@ -0,0 +1,24 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched + <192, + 128, + MaskMode::kMultiItemScoring, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultAttention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched + <192, + 128, + MaskMode::kMultiItemScoring, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultAttention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_0.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..a8ff81056119dcb0fe6a8e23661212bcbb9df29c --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_0.cu @@ -0,0 +1,24 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched + <192, + 128, + MaskMode::kNone, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultAttention>(RaggedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched + <192, + 128, + MaskMode::kNone, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultAttention>(RaggedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_1.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..298cab570997fa4520c656188979e48efa6bb65b --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_1.cu @@ -0,0 +1,24 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched + <192, + 128, + MaskMode::kCausal, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultAttention>(RaggedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched + <192, + 128, + MaskMode::kCausal, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultAttention>(RaggedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_2.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..40ef2acdb67e3f7748afb7c8dea728bbc407fd3d --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_2.cu @@ -0,0 +1,24 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched + <192, + 128, + MaskMode::kCustom, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultAttention>(RaggedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched + <192, + 128, + MaskMode::kCustom, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultAttention>(RaggedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_3.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..f50854c38ab7a25526c070ff8eb10550151ed35b --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_3.cu @@ -0,0 +1,24 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched + <192, + 128, + MaskMode::kMultiItemScoring, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultAttention>(RaggedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched + <192, + 128, + MaskMode::kMultiItemScoring, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultAttention>(RaggedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_sm90.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_sm90.cu new file mode 100644 index 0000000000000000000000000000000000000000..159ee3f1eda3f47769ec14f8fac16f35fa9bda05 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_sm90.cu @@ -0,0 +1,260 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include +#include + +#include "batch_prefill_sm90_config.inc" +#include "pytorch_conversion_utils.h" +#include "pytorch_extension_utils.h" + +namespace flashinfer { + +template +cudaError_t BatchPrefillWithRaggedKVCacheDispatched(Params& params, bool enable_pdl, + cudaStream_t stream); + +template +cudaError_t BatchPrefillWithPagedKVCacheDispatched(Params& params, bool enable_pdl, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +at::Tensor BatchPrefillWithKVCacheSM90Plan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, + int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, + int64_t head_dim_vo, bool causal) { + size_t float_workspace_size_in_bytes = + float_workspace_buffer.size(0) * float_workspace_buffer.element_size(); + size_t int_workspace_size_in_bytes = + int_workspace_buffer.size(0) * int_workspace_buffer.element_size(); + + flashinfer::PrefillPlanSM90Info plan_info; + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + cudaError_t status = + PrefillSM90Plan(float_workspace_buffer.data_ptr(), float_workspace_size_in_bytes, + int_workspace_buffer.data_ptr(), page_locked_int_workspace_buffer.data_ptr(), + int_workspace_size_in_bytes, plan_info, qo_indptr.data_ptr(), + kv_indptr.data_ptr(), kv_len_arr.data_ptr(), total_num_rows, + batch_size, num_qo_heads, num_kv_heads, head_dim_qk, head_dim_vo, page_size, + causal, enable_cuda_graph, /*sizeof_dtype_o=*/2, stream); + + TORCH_CHECK(status == cudaSuccess, + "PrefillSM90Plan failed with error: ", cudaGetErrorString(status)); + + return vec_to_tensor(plan_info.ToVector()); +} + +void BatchPrefillWithRaggedKVCacheSM90Run( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + PrefillPlanSM90Info plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == q.size(0), lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == q.size(1), lse.size(1), q.size(1)); + } + + void* float_buffer_ptr = float_workspace_buffer.data_ptr(); + void* int_buffer_ptr = int_workspace_buffer.data_ptr(); + + int64_t head_dim_qk = q.size(2); + int64_t head_dim_vo = v.size(2); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + QKVLayout kv_layout = static_cast(layout); + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + const MaskMode mask_mode = static_cast(mask_mode_code); + bool use_swa = window_left != -1; + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, USE_SLIDING_WINDOW, + USE_LOGITS_SOFT_CAP, AttentionVariant, RaggedParams, PagedParams, [&] { + RaggedParams params; + + params.q_ptr = static_cast(q.data_ptr()); + params.k_ptr = static_cast(k.data_ptr()); + params.v_ptr = static_cast(v.data_ptr()); + params.o_ptr = static_cast(o.data_ptr()); + params.lse_ptr = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.q_stride_n = q.stride(0); + params.q_stride_h = q.stride(1); + params.o_stride_n = o.stride(0); + params.o_stride_h = o.stride(1); + if (kv_layout == QKVLayout::kNHD) { + params.k_stride_n = k.stride(0); + params.k_stride_h = k.stride(1); + params.v_stride_n = v.stride(0); + params.v_stride_h = v.stride(1); + } else { + params.k_stride_h = k.stride(0); + params.k_stride_n = k.stride(1); + params.v_stride_h = v.stride(0); + params.v_stride_n = v.stride(1); + } + params.nnz_qo = q.size(0); + params.nnz_kv = k.size(0); + params.num_qo_heads = q.size(1); + params.num_kv_heads = k.size(1); + params.group_size = params.num_qo_heads / params.num_kv_heads; + params.window_left = window_left; + params.causal = mask_mode_code == 1; + params.qo_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_tile_indices_offset); + params.qo_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_indptr_offset); + params.kv_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_indptr_offset); + params.qo_lens = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_len_offset); + params.kv_lens = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_len_offset); + params.head_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.head_indices_offset); + params.work_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.work_indptr_offset); + params.batch_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.batch_indices_offset); + + ADDITIONAL_PARAMS_SETTER + + bool same_schedule_for_all_heads = plan_info.same_schedule_for_all_heads; + DISPATCH_BOOL(same_schedule_for_all_heads, SAME_SCHEDULER_FOR_ALL_HEADS, [&] { + cudaError_t status = BatchPrefillWithRaggedKVCacheDispatched< + HEAD_DIM_QK, HEAD_DIM_VO, MASK_MODE, USE_SLIDING_WINDOW, SAME_SCHEDULER_FOR_ALL_HEADS, + AttentionVariant>(params, enable_pdl, stream); + TORCH_CHECK(status == cudaSuccess, + "BatchPrefillWithRaggedKVCacheSM90Run failed with error: ", + cudaGetErrorString(status)); + return true; + }); + }); +} + +void BatchPrefillWithPagedKVCacheSM90Run( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, + at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + PrefillPlanSM90Info plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == q.size(0), lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == q.size(1), lse.size(1), q.size(1)); + } + QKVLayout kv_layout = static_cast(layout); + int64_t num_kv_heads, page_size; + int64_t head_dim_qk = q.size(2); + int64_t head_dim_vo = paged_v_cache.size(3); + if (kv_layout == QKVLayout::kHND) { + num_kv_heads = paged_k_cache.size(1); + page_size = paged_k_cache.size(2); + } else { + page_size = paged_k_cache.size(1); + num_kv_heads = paged_k_cache.size(2); + } + + void* float_buffer_ptr = float_workspace_buffer.data_ptr(); + void* int_buffer_ptr = int_workspace_buffer.data_ptr(); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = paged_k_cache.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + const MaskMode mask_mode = static_cast(mask_mode_code); + bool use_swa = window_left != -1; + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, USE_SLIDING_WINDOW, + USE_LOGITS_SOFT_CAP, AttentionVariant, RaggedParams, PagedParams, [&] { + PagedParams params; + + params.q_ptr = static_cast(q.data_ptr()); + params.k_ptr = static_cast(paged_k_cache.data_ptr()); + params.v_ptr = static_cast(paged_v_cache.data_ptr()); + params.o_ptr = static_cast(o.data_ptr()); + params.lse_ptr = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.q_stride_n = q.stride(0); + params.q_stride_h = q.stride(1); + params.o_stride_n = o.stride(0); + params.o_stride_h = o.stride(1); + if (kv_layout == QKVLayout::kNHD) { + // (num_pages, page_size, num_heads, head_dim) + params.k_stride_n = paged_k_cache.stride(1); + params.k_stride_h = paged_k_cache.stride(2); + params.v_stride_n = paged_v_cache.stride(1); + params.v_stride_h = paged_v_cache.stride(2); + } else { + // (num_pages, num_heads, page_size, head_dim) + params.k_stride_h = paged_k_cache.stride(1); + params.k_stride_n = paged_k_cache.stride(2); + params.v_stride_h = paged_v_cache.stride(1); + params.v_stride_n = paged_v_cache.stride(2); + } + params.nnz_qo = q.size(0); + params.num_qo_heads = q.size(1); + params.num_kv_heads = num_kv_heads; + params.group_size = params.num_qo_heads / num_kv_heads; + params.page_size = page_size; + params.window_left = window_left; + params.causal = mask_mode_code == 1; + params.qo_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_tile_indices_offset); + params.qo_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_indptr_offset); + params.kv_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_indptr_offset); + params.qo_lens = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_len_offset); + params.kv_lens = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_len_offset); + params.head_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.head_indices_offset); + params.work_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.work_indptr_offset); + params.batch_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.batch_indices_offset); + params.kv_indices = static_cast(paged_kv_indices.data_ptr()); + + ADDITIONAL_PARAMS_SETTER + + bool same_schedule_for_all_heads = plan_info.same_schedule_for_all_heads; + DISPATCH_BOOL(same_schedule_for_all_heads, SAME_SCHEDULER_FOR_ALL_HEADS, [&] { + cudaError_t status = BatchPrefillWithPagedKVCacheDispatched< + HEAD_DIM_QK, HEAD_DIM_VO, MASK_MODE, USE_SLIDING_WINDOW, SAME_SCHEDULER_FOR_ALL_HEADS, + AttentionVariant>(params, enable_pdl, stream); + TORCH_CHECK(status == cudaSuccess, + "BatchPrefillWithPagedKVCacheSM90Run failed with error: ", + cudaGetErrorString(status)); + return true; + }); + }); +} diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_sm90_config.inc b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_sm90_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..f90116827437b4277434dfe2c2ee529220b4e661 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_sm90_config.inc @@ -0,0 +1,135 @@ +#pragma once +#include +#include +#include +#include +#include +#include "pytorch_extension_utils.h" + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_prefix_len_ptr, std::optional maybe_token_pos_in_items_ptr, std::optional maybe_max_item_len_ptr, double logits_soft_cap, double sm_scale, int64_t token_pos_in_items_len +#define ADDITIONAL_PARAMS_SETTER params.additional_params.maybe_prefix_len_ptr = maybe_prefix_len_ptr ? static_cast(maybe_prefix_len_ptr->data_ptr()): nullptr; \ +params.additional_params.maybe_token_pos_in_items_ptr = maybe_token_pos_in_items_ptr ? static_cast(maybe_token_pos_in_items_ptr->data_ptr()): nullptr; \ +params.additional_params.maybe_max_item_len_ptr = maybe_max_item_len_ptr ? static_cast(maybe_max_item_len_ptr->data_ptr()): nullptr; \ +params.additional_params.logits_soft_cap = logits_soft_cap; \ +params.additional_params.sm_scale = sm_scale; \ +params.additional_params.token_pos_in_items_len = token_pos_in_items_len; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, RaggedParams, PagedParams, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { using AttentionVariant = DefaultAttention; __VA_ARGS__();}) + +using namespace flashinfer; + +using DTypeQ = cutlass_dtype_t; +using DTypeKV = cutlass_dtype_t; +using DTypeO = cutlass_dtype_t; +using IdType = cutlass_dtype_t; + +constexpr int HEAD_DIM_QK = 192; +constexpr int HEAD_DIM_VO = 128; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto USE_SLIDING_WINDOW = false; + +struct RaggedParams { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + // The QKV matrices. + DTypeQ* q_ptr; + DTypeKV* k_ptr; + DTypeKV* v_ptr; + DTypeO* o_ptr; + float* lse_ptr; + + IdType* qo_tile_indices; + IdType* qo_indptr; + IdType* kv_indptr; + IdType* qo_lens; + IdType* kv_lens; + IdType* head_indices; + IdType* work_indptr; + IdType* batch_indices; + + struct AdditionalParams { + uint32_t* maybe_prefix_len_ptr; +uint16_t* maybe_token_pos_in_items_ptr; +uint16_t* maybe_max_item_len_ptr; +double logits_soft_cap; +double sm_scale; +int64_t token_pos_in_items_len; + + } additional_params; + + int64_t q_stride_n; + int64_t k_stride_n; + int64_t v_stride_n; + int64_t o_stride_n; + int64_t q_stride_h; + int64_t k_stride_h; + int64_t v_stride_h; + int64_t o_stride_h; + int64_t nnz_qo; + int64_t nnz_kv; + + int head_dim; + int num_qo_heads; + int num_kv_heads; + int group_size; + int window_left; + + bool causal; +}; + +struct PagedParams { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + // The QKV matrices. + DTypeQ* q_ptr; + DTypeKV* k_ptr; + DTypeKV* v_ptr; + DTypeO* o_ptr; + float* lse_ptr; + + IdType* qo_tile_indices; + IdType* qo_indptr; + IdType* kv_indptr; + IdType* kv_indices; + IdType* qo_lens; + IdType* kv_lens; + IdType* head_indices; + IdType* work_indptr; + IdType* batch_indices; + + struct AdditionalParams { + uint32_t* maybe_prefix_len_ptr; +uint16_t* maybe_token_pos_in_items_ptr; +uint16_t* maybe_max_item_len_ptr; +double logits_soft_cap; +double sm_scale; +int64_t token_pos_in_items_len; + + } additional_params; + + int64_t q_stride_n; + int64_t k_stride_n; + int64_t v_stride_n; + int64_t o_stride_n; + int64_t q_stride_h; + int64_t k_stride_h; + int64_t v_stride_h; + int64_t o_stride_h; + int64_t nnz_qo; + + int head_dim; + int num_qo_heads; + int num_kv_heads; + int group_size; + int page_size; + int window_left; + + bool causal; +}; + +#include \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_sm90_jit_pybind.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_sm90_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..b1edd582549e124c776f2cf9b4503f4bf60d16ae --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_sm90_jit_pybind.cu @@ -0,0 +1,46 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "batch_prefill_sm90_config.inc" +#include "pytorch_extension_utils.h" + +at::Tensor BatchPrefillWithKVCacheSM90Plan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, + int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, + int64_t head_dim_vo, bool causal); + +void BatchPrefillWithRaggedKVCacheSM90Run( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +void BatchPrefillWithPagedKVCacheSM90Run( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, + at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Batch-request prefill attention with KV-Cache plan + m.def("plan", BatchPrefillWithKVCacheSM90Plan); + // Batch-request prefill attention with KV-Cache operator + m.def("ragged_run", BatchPrefillWithRaggedKVCacheSM90Run); + // Batch-request prefill attention with KV-Cache operator + m.def("paged_run", BatchPrefillWithPagedKVCacheSM90Run); +} diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill.cu new file mode 100644 index 0000000000000000000000000000000000000000..a51fc7f56a0af7482bd5d2a9d3f672854df69dfd --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill.cu @@ -0,0 +1,330 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include +#include + +#include "batch_prefill_config.inc" +#include "pytorch_conversion_utils.h" +#include "pytorch_extension_utils.h" + +namespace flashinfer { + +template +cudaError_t BatchPrefillWithPagedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v, + float* tmp_s, bool enable_pdl, + cudaStream_t stream); + +template +cudaError_t BatchPrefillWithRaggedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v, + float* tmp_s, bool enable_pdl, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +at::Tensor BatchPrefillWithKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, + int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, + int64_t head_dim_vo, bool causal) { + size_t float_workspace_size_in_bytes = + float_workspace_buffer.size(0) * float_workspace_buffer.element_size(); + size_t int_workspace_size_in_bytes = + int_workspace_buffer.size(0) * int_workspace_buffer.element_size(); + + PrefillPlanInfo plan_info; + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + cudaError_t status = PrefillPlan( + float_workspace_buffer.data_ptr(), float_workspace_size_in_bytes, + int_workspace_buffer.data_ptr(), page_locked_int_workspace_buffer.data_ptr(), + int_workspace_size_in_bytes, plan_info, qo_indptr.data_ptr(), + kv_indptr.data_ptr(), total_num_rows, batch_size, num_qo_heads, num_kv_heads, + head_dim_qk, head_dim_vo, page_size, enable_cuda_graph, /*sizeof_dtype_o=*/2, stream); + + TORCH_CHECK(status == cudaSuccess, + "Failed to plan prefill with error: ", cudaGetErrorString(status)); + + return vec_to_tensor(plan_info.ToVector()); +} + +void BatchPrefillWithRaggedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor k, at::Tensor v, + at::Tensor qo_indptr, at::Tensor kv_indptr, at::Tensor o, + std::optional maybe_lse, int64_t mask_mode_code, + int64_t layout, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + PrefillPlanInfo plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + QKVLayout kv_layout = static_cast(layout); + + int64_t num_qo_heads = q.size(1); + int64_t head_dim_qk = q.size(2); + int64_t num_kv_heads = (kv_layout == QKVLayout::kNHD) ? k.size(1) : k.size(0); + uint32_t q_stride_n = q.stride(0), q_stride_h = q.stride(1), k_stride_n, k_stride_h, v_stride_n, + v_stride_h; + if (kv_layout == QKVLayout::kNHD) { + k_stride_n = k.stride(0); + k_stride_h = k.stride(1); + v_stride_n = v.stride(0); + v_stride_h = v.stride(1); + } else { + k_stride_h = k.stride(0); + k_stride_n = k.stride(1); + v_stride_h = v.stride(0); + v_stride_n = v.stride(1); + } + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == q.size(0), lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == q.size(1), lse.size(1), q.size(1)); + } + + void* float_buffer_ptr = float_workspace_buffer.data_ptr(); + void* int_buffer_ptr = int_workspace_buffer.data_ptr(); + + const MaskMode mask_mode = static_cast(mask_mode_code); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, + RaggedParams, PagedParams, [&] { + RaggedParams params; + + params.q = static_cast(q.data_ptr()); + params.k = static_cast(k.data_ptr()); + params.v = static_cast(v.data_ptr()); + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.q_indptr = static_cast(qo_indptr.data_ptr()); + params.kv_indptr = static_cast(kv_indptr.data_ptr()); + params.num_qo_heads = num_qo_heads; + params.num_kv_heads = num_kv_heads; + params.group_size = uint_fastdiv(num_qo_heads / num_kv_heads); + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.k_stride_n = k_stride_n; + params.k_stride_h = k_stride_h; + params.v_stride_n = v_stride_n; + params.v_stride_h = v_stride_h; + params.window_left = window_left; + + params.request_indices = nullptr; + params.qo_tile_indices = nullptr; + params.kv_tile_indices = nullptr; + params.merge_indptr = nullptr; + params.o_indptr = nullptr; + params.kv_chunk_size_ptr = nullptr; + params.block_valid_mask = nullptr; + params.total_num_rows = nullptr; + params.max_total_num_rows = 0; + params.padded_batch_size = 0; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + DTypeO* tmp_v = nullptr; + float* tmp_s = nullptr; + + params.request_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.request_indices_offset); + params.qo_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_tile_indices_offset); + params.kv_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_tile_indices_offset); + params.o_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.o_indptr_offset); + params.kv_chunk_size_ptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_chunk_size_ptr_offset); + if (plan_info.split_kv) { + params.merge_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.merge_indptr_offset); + tmp_v = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.v_offset); + tmp_s = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.s_offset); + if (plan_info.enable_cuda_graph) { + params.block_valid_mask = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.block_valid_mask_offset); + } + } + params.padded_batch_size = plan_info.padded_batch_size; + params.max_total_num_rows = plan_info.total_num_rows; + if (plan_info.enable_cuda_graph) { + params.total_num_rows = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.total_num_rows_offset); + } + + cudaError_t status = cudaSuccess; + + DISPATCH_CTA_TILE_Q(plan_info.cta_tile_q, CTA_TILE_Q, { + status = flashinfer::BatchPrefillWithRaggedKVCacheDispatched< + CTA_TILE_Q, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + /*use_fp16_qk_reduction=*/USE_FP16_QK_REDUCTION, MASK_MODE, AttentionVariant, + RaggedParams>(params, tmp_v, tmp_s, enable_pdl, stream); + }); + + TORCH_CHECK(status == cudaSuccess, "BatchPrefillWithRaggedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); +} + +void BatchPrefillWithPagedKVCacheRun( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, + at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + PrefillPlanInfo plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + QKVLayout kv_layout = static_cast(layout); + auto device = q.device(); + int64_t batch_size = paged_kv_indptr.size(0) - 1; + int64_t num_qo_heads = q.size(1); + int64_t num_kv_heads, page_size; + uint32_t head_dim_qk = q.size(2); + if (kv_layout == QKVLayout::kHND) { + num_kv_heads = paged_k_cache.size(1); + page_size = paged_k_cache.size(2); + } else { + page_size = paged_k_cache.size(1); + num_kv_heads = paged_k_cache.size(2); + } + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == q.size(0), lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == q.size(1), lse.size(1), q.size(1)); + } + + void* float_buffer_ptr = static_cast(float_workspace_buffer.data_ptr()); + void* int_buffer_ptr = static_cast(int_workspace_buffer.data_ptr()); + + const MaskMode mask_mode = static_cast(mask_mode_code); + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = paged_k_cache.scalar_type(); + + // get q_stride_n and q_stride_h + const auto q_stride_n = q.stride(0); + const auto q_stride_h = q.stride(1); + + // get kv_cache_strides + const int64_t* kv_cache_strides = nullptr; + auto k_strides = paged_k_cache.strides(); + auto v_strides = paged_v_cache.strides(); + TORCH_CHECK(k_strides == v_strides, "k/v strides must be identical"); + kv_cache_strides = k_strides.data(); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, + RaggedParams, PagedParams, [&] { + PagedParams params; + + params.q = static_cast(q.data_ptr()); + paged_kv_t paged_kv( + num_kv_heads, page_size, HEAD_DIM_VO, batch_size, kv_layout, + static_cast(paged_k_cache.data_ptr()), + static_cast(paged_v_cache.data_ptr()), kv_cache_strides, + static_cast(paged_kv_indices.data_ptr()), + static_cast(paged_kv_indptr.data_ptr()), + static_cast(paged_kv_last_page_len.data_ptr())); + params.paged_kv = paged_kv; + params.q_indptr = static_cast(qo_indptr.data_ptr()); + params.o = static_cast(o.data_ptr()); + + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.num_qo_heads = num_qo_heads; + params.group_size = uint_fastdiv(num_qo_heads / paged_kv.num_heads); + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.window_left = window_left; + + params.request_indices = nullptr; + params.qo_tile_indices = nullptr; + params.kv_tile_indices = nullptr; + params.merge_indptr = nullptr; + params.o_indptr = nullptr; + params.kv_chunk_size_ptr = nullptr; + params.block_valid_mask = nullptr; + params.total_num_rows = nullptr; + params.max_total_num_rows = 0; + params.padded_batch_size = 0; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + DTypeO* tmp_v = nullptr; + float* tmp_s = nullptr; + + params.request_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.request_indices_offset); + params.qo_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_tile_indices_offset); + params.kv_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_tile_indices_offset); + params.o_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.o_indptr_offset); + params.kv_chunk_size_ptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_chunk_size_ptr_offset); + if (plan_info.split_kv) { + params.merge_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.merge_indptr_offset); + tmp_v = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.v_offset); + tmp_s = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.s_offset); + if (plan_info.enable_cuda_graph) { + params.block_valid_mask = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.block_valid_mask_offset); + } + } + params.padded_batch_size = plan_info.padded_batch_size; + params.max_total_num_rows = plan_info.total_num_rows; + if (plan_info.enable_cuda_graph) { + params.total_num_rows = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.total_num_rows_offset); + } + + cudaError_t status = cudaSuccess; + + DISPATCH_CTA_TILE_Q(plan_info.cta_tile_q, CTA_TILE_Q, { + status = flashinfer::BatchPrefillWithPagedKVCacheDispatched< + CTA_TILE_Q, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + /*use_fp16_qk_reduction=*/USE_FP16_QK_REDUCTION, MASK_MODE, AttentionVariant, + PagedParams>(params, tmp_v, tmp_s, enable_pdl, stream); + }); + + TORCH_CHECK(status == cudaSuccess, "BatchPrefillWithPagedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); +} diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_config.inc b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..73e9ce745c82041e9b38fa33f87b5692d45175b4 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_config.inc @@ -0,0 +1,153 @@ +#pragma once +#include +#include +#include +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_custom_mask, std::optional maybe_mask_indptr, std::optional maybe_alibi_slopes, std::optional maybe_prefix_len_ptr, std::optional maybe_token_pos_in_items_ptr, std::optional maybe_max_item_len_ptr, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta, int64_t token_pos_in_items_len +#define ADDITIONAL_PARAMS_SETTER params.maybe_custom_mask = maybe_custom_mask ? static_cast(maybe_custom_mask->data_ptr()): nullptr; \ +params.maybe_mask_indptr = maybe_mask_indptr ? static_cast(maybe_mask_indptr->data_ptr()): nullptr; \ +params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.maybe_prefix_len_ptr = maybe_prefix_len_ptr ? static_cast(maybe_prefix_len_ptr->data_ptr()): nullptr; \ +params.maybe_token_pos_in_items_ptr = maybe_token_pos_in_items_ptr ? static_cast(maybe_token_pos_in_items_ptr->data_ptr()): nullptr; \ +params.maybe_max_item_len_ptr = maybe_max_item_len_ptr ? static_cast(maybe_max_item_len_ptr->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; \ +params.token_pos_in_items_len = token_pos_in_items_len; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, RaggedParams, PagedParams, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { \ + constexpr auto use_custom_mask = MASK_MODE == MaskMode::kCustom; \ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ + }) + +using namespace flashinfer; + +using DTypeQ = nv_bfloat16; +using DTypeKV = nv_bfloat16; +using DTypeO = nv_bfloat16; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 256; +constexpr int HEAD_DIM_VO = 256; +constexpr bool USE_FP16_QK_REDUCTION = false; +constexpr auto USE_LOGITS_SOFT_CAP = true; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = true; + + +struct RaggedParams { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + DTypeQ* q; + DTypeKV* k; + DTypeKV* v; + IdType* q_indptr; + IdType* kv_indptr; + DTypeO* o; + float* lse; + uint_fastdiv group_size; + + uint8_t* maybe_custom_mask; +int32_t* maybe_mask_indptr; +float* maybe_alibi_slopes; +uint32_t* maybe_prefix_len_ptr; +uint16_t* maybe_token_pos_in_items_ptr; +uint16_t* maybe_max_item_len_ptr; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; +int64_t token_pos_in_items_len; + + uint32_t num_qo_heads; + uint32_t num_kv_heads; + uint32_t q_stride_n; + uint32_t q_stride_h; + uint32_t k_stride_n; + uint32_t k_stride_h; + uint32_t v_stride_n; + uint32_t v_stride_h; + int32_t window_left; + + IdType* request_indices; + IdType* qo_tile_indices; + IdType* kv_tile_indices; + IdType* merge_indptr; + IdType* o_indptr; + IdType* kv_chunk_size_ptr; + bool* block_valid_mask; + uint32_t max_total_num_rows; + uint32_t* total_num_rows; + uint32_t padded_batch_size; + bool partition_kv; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { + return q_indptr[batch_idx + 1] - q_indptr[batch_idx]; + } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return kv_indptr[batch_idx + 1] - kv_indptr[batch_idx]; + } +}; + +struct PagedParams { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + DTypeQ* q; + paged_kv_t paged_kv; + IdType* q_indptr; + DTypeO* o; + float* lse; + uint_fastdiv group_size; + + uint8_t* maybe_custom_mask; +int32_t* maybe_mask_indptr; +float* maybe_alibi_slopes; +uint32_t* maybe_prefix_len_ptr; +uint16_t* maybe_token_pos_in_items_ptr; +uint16_t* maybe_max_item_len_ptr; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; +int64_t token_pos_in_items_len; + + uint32_t num_qo_heads; + IdType q_stride_n; + IdType q_stride_h; + int32_t window_left; + + IdType* request_indices; + IdType* qo_tile_indices; + IdType* kv_tile_indices; + IdType* merge_indptr; + IdType* o_indptr; + bool* block_valid_mask; + IdType* kv_chunk_size_ptr; + uint32_t max_total_num_rows; + uint32_t* total_num_rows; + uint32_t padded_batch_size; + bool partition_kv; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { + return q_indptr[batch_idx + 1] - q_indptr[batch_idx]; + } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return paged_kv.get_length(batch_idx); + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_jit_pybind.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..5421ab1cf910fdfbc867c182c8e22f23e44749ae --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_jit_pybind.cu @@ -0,0 +1,48 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "batch_prefill_config.inc" +#include "pytorch_extension_utils.h" + +at::Tensor BatchPrefillWithKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, + int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, + int64_t head_dim_vo, bool causal); + +void BatchPrefillWithRaggedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor k, at::Tensor v, + at::Tensor qo_indptr, at::Tensor kv_indptr, at::Tensor o, + std::optional maybe_lse, int64_t mask_mode_code, + int64_t layout, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +void BatchPrefillWithPagedKVCacheRun( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, + at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Batch-request prefill attention with KV-Cache plan + m.def("plan", BatchPrefillWithKVCachePlan); + // Batch-request prefill attention with KV-Cache operator + m.def("ragged_run", BatchPrefillWithRaggedKVCacheRun); + // Batch-request prefill attention with KV-Cache operator + m.def("paged_run", BatchPrefillWithPagedKVCacheRun); +} diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_paged_kernel_mask_0.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_paged_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..264348261256ffc911b50c9fd7237441faec122f --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_paged_kernel_mask_0.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kNone == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 256, 256, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 256, 256, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 256, 256, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_paged_kernel_mask_1.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_paged_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..9937900da6dca1744966105d9de12f7321ae00e4 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_paged_kernel_mask_1.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCausal == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 256, 256, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 256, 256, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 256, 256, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_paged_kernel_mask_2.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_paged_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..f61a91898a844ffbce56d04112d87a9e078243d8 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_paged_kernel_mask_2.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCustom == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 256, 256, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 256, 256, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 256, 256, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_paged_kernel_mask_3.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_paged_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..e258a720abb5b0a909466b307c58f19ddd43c157 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_paged_kernel_mask_3.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kMultiItemScoring == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 256, 256, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 256, 256, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 256, 256, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_ragged_kernel_mask_0.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_ragged_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..ff099d149943dfad67e18d6b600d3f5781bf5d02 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_ragged_kernel_mask_0.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kNone == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 256, 256, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 256, 256, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 256, 256, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_ragged_kernel_mask_1.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_ragged_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..838dcf44be11d68bb2c48ddbc70c8b42c62d4f93 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_ragged_kernel_mask_1.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCausal == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 256, 256, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 256, 256, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 256, 256, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_ragged_kernel_mask_2.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_ragged_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..9735995235347e9cf67269a5d532eaf67df02b14 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_ragged_kernel_mask_2.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCustom == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 256, 256, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 256, 256, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 256, 256, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_ragged_kernel_mask_3.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_ragged_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..255af5bee1786391bc5be516b8f8c9c051133a41 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_ragged_kernel_mask_3.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kMultiItemScoring == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 256, 256, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 256, 256, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 256, 256, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_0.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..700eb2e367c50b3f4e1bf963849d60e099c03a99 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_0.cu @@ -0,0 +1,24 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched + <256, + 256, + MaskMode::kNone, + /*USE_SLIDING_WINDOW=*/true, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultAttention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched + <256, + 256, + MaskMode::kNone, + /*USE_SLIDING_WINDOW=*/true, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultAttention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_1.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..ae7ae1a87dc15bcbb01569c0ad7fa25ef0f3e8f7 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_1.cu @@ -0,0 +1,24 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched + <256, + 256, + MaskMode::kCausal, + /*USE_SLIDING_WINDOW=*/true, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultAttention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched + <256, + 256, + MaskMode::kCausal, + /*USE_SLIDING_WINDOW=*/true, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultAttention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_2.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..d2a680686e6ea2a4611481277d89a2fc3608d9f4 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_2.cu @@ -0,0 +1,24 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched + <256, + 256, + MaskMode::kCustom, + /*USE_SLIDING_WINDOW=*/true, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultAttention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched + <256, + 256, + MaskMode::kCustom, + /*USE_SLIDING_WINDOW=*/true, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultAttention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_3.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..4a0f46b6c308a6fd4dcf12ae2d4be11fa5bd89eb --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_3.cu @@ -0,0 +1,24 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched + <256, + 256, + MaskMode::kMultiItemScoring, + /*USE_SLIDING_WINDOW=*/true, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultAttention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched + <256, + 256, + MaskMode::kMultiItemScoring, + /*USE_SLIDING_WINDOW=*/true, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultAttention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_0.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..afd17ddcf65342899385c0094d7187f93b92756d --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_0.cu @@ -0,0 +1,24 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched + <256, + 256, + MaskMode::kNone, + /*USE_SLIDING_WINDOW=*/true, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultAttention>(RaggedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched + <256, + 256, + MaskMode::kNone, + /*USE_SLIDING_WINDOW=*/true, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultAttention>(RaggedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_1.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..c901f44d76128501cc772b356f5217211199bb50 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_1.cu @@ -0,0 +1,24 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched + <256, + 256, + MaskMode::kCausal, + /*USE_SLIDING_WINDOW=*/true, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultAttention>(RaggedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched + <256, + 256, + MaskMode::kCausal, + /*USE_SLIDING_WINDOW=*/true, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultAttention>(RaggedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_2.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..00b399ef1c529610375703d79f261333f1c1c386 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_2.cu @@ -0,0 +1,24 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched + <256, + 256, + MaskMode::kCustom, + /*USE_SLIDING_WINDOW=*/true, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultAttention>(RaggedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched + <256, + 256, + MaskMode::kCustom, + /*USE_SLIDING_WINDOW=*/true, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultAttention>(RaggedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_3.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..d9da788a703f3c97c596380ea9eccb4427ed228a --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_3.cu @@ -0,0 +1,24 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched + <256, + 256, + MaskMode::kMultiItemScoring, + /*USE_SLIDING_WINDOW=*/true, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultAttention>(RaggedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched + <256, + 256, + MaskMode::kMultiItemScoring, + /*USE_SLIDING_WINDOW=*/true, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultAttention>(RaggedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_sm90.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_sm90.cu new file mode 100644 index 0000000000000000000000000000000000000000..159ee3f1eda3f47769ec14f8fac16f35fa9bda05 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_sm90.cu @@ -0,0 +1,260 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include +#include + +#include "batch_prefill_sm90_config.inc" +#include "pytorch_conversion_utils.h" +#include "pytorch_extension_utils.h" + +namespace flashinfer { + +template +cudaError_t BatchPrefillWithRaggedKVCacheDispatched(Params& params, bool enable_pdl, + cudaStream_t stream); + +template +cudaError_t BatchPrefillWithPagedKVCacheDispatched(Params& params, bool enable_pdl, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +at::Tensor BatchPrefillWithKVCacheSM90Plan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, + int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, + int64_t head_dim_vo, bool causal) { + size_t float_workspace_size_in_bytes = + float_workspace_buffer.size(0) * float_workspace_buffer.element_size(); + size_t int_workspace_size_in_bytes = + int_workspace_buffer.size(0) * int_workspace_buffer.element_size(); + + flashinfer::PrefillPlanSM90Info plan_info; + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + cudaError_t status = + PrefillSM90Plan(float_workspace_buffer.data_ptr(), float_workspace_size_in_bytes, + int_workspace_buffer.data_ptr(), page_locked_int_workspace_buffer.data_ptr(), + int_workspace_size_in_bytes, plan_info, qo_indptr.data_ptr(), + kv_indptr.data_ptr(), kv_len_arr.data_ptr(), total_num_rows, + batch_size, num_qo_heads, num_kv_heads, head_dim_qk, head_dim_vo, page_size, + causal, enable_cuda_graph, /*sizeof_dtype_o=*/2, stream); + + TORCH_CHECK(status == cudaSuccess, + "PrefillSM90Plan failed with error: ", cudaGetErrorString(status)); + + return vec_to_tensor(plan_info.ToVector()); +} + +void BatchPrefillWithRaggedKVCacheSM90Run( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + PrefillPlanSM90Info plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == q.size(0), lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == q.size(1), lse.size(1), q.size(1)); + } + + void* float_buffer_ptr = float_workspace_buffer.data_ptr(); + void* int_buffer_ptr = int_workspace_buffer.data_ptr(); + + int64_t head_dim_qk = q.size(2); + int64_t head_dim_vo = v.size(2); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + QKVLayout kv_layout = static_cast(layout); + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + const MaskMode mask_mode = static_cast(mask_mode_code); + bool use_swa = window_left != -1; + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, USE_SLIDING_WINDOW, + USE_LOGITS_SOFT_CAP, AttentionVariant, RaggedParams, PagedParams, [&] { + RaggedParams params; + + params.q_ptr = static_cast(q.data_ptr()); + params.k_ptr = static_cast(k.data_ptr()); + params.v_ptr = static_cast(v.data_ptr()); + params.o_ptr = static_cast(o.data_ptr()); + params.lse_ptr = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.q_stride_n = q.stride(0); + params.q_stride_h = q.stride(1); + params.o_stride_n = o.stride(0); + params.o_stride_h = o.stride(1); + if (kv_layout == QKVLayout::kNHD) { + params.k_stride_n = k.stride(0); + params.k_stride_h = k.stride(1); + params.v_stride_n = v.stride(0); + params.v_stride_h = v.stride(1); + } else { + params.k_stride_h = k.stride(0); + params.k_stride_n = k.stride(1); + params.v_stride_h = v.stride(0); + params.v_stride_n = v.stride(1); + } + params.nnz_qo = q.size(0); + params.nnz_kv = k.size(0); + params.num_qo_heads = q.size(1); + params.num_kv_heads = k.size(1); + params.group_size = params.num_qo_heads / params.num_kv_heads; + params.window_left = window_left; + params.causal = mask_mode_code == 1; + params.qo_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_tile_indices_offset); + params.qo_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_indptr_offset); + params.kv_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_indptr_offset); + params.qo_lens = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_len_offset); + params.kv_lens = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_len_offset); + params.head_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.head_indices_offset); + params.work_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.work_indptr_offset); + params.batch_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.batch_indices_offset); + + ADDITIONAL_PARAMS_SETTER + + bool same_schedule_for_all_heads = plan_info.same_schedule_for_all_heads; + DISPATCH_BOOL(same_schedule_for_all_heads, SAME_SCHEDULER_FOR_ALL_HEADS, [&] { + cudaError_t status = BatchPrefillWithRaggedKVCacheDispatched< + HEAD_DIM_QK, HEAD_DIM_VO, MASK_MODE, USE_SLIDING_WINDOW, SAME_SCHEDULER_FOR_ALL_HEADS, + AttentionVariant>(params, enable_pdl, stream); + TORCH_CHECK(status == cudaSuccess, + "BatchPrefillWithRaggedKVCacheSM90Run failed with error: ", + cudaGetErrorString(status)); + return true; + }); + }); +} + +void BatchPrefillWithPagedKVCacheSM90Run( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, + at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + PrefillPlanSM90Info plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == q.size(0), lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == q.size(1), lse.size(1), q.size(1)); + } + QKVLayout kv_layout = static_cast(layout); + int64_t num_kv_heads, page_size; + int64_t head_dim_qk = q.size(2); + int64_t head_dim_vo = paged_v_cache.size(3); + if (kv_layout == QKVLayout::kHND) { + num_kv_heads = paged_k_cache.size(1); + page_size = paged_k_cache.size(2); + } else { + page_size = paged_k_cache.size(1); + num_kv_heads = paged_k_cache.size(2); + } + + void* float_buffer_ptr = float_workspace_buffer.data_ptr(); + void* int_buffer_ptr = int_workspace_buffer.data_ptr(); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = paged_k_cache.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + const MaskMode mask_mode = static_cast(mask_mode_code); + bool use_swa = window_left != -1; + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, USE_SLIDING_WINDOW, + USE_LOGITS_SOFT_CAP, AttentionVariant, RaggedParams, PagedParams, [&] { + PagedParams params; + + params.q_ptr = static_cast(q.data_ptr()); + params.k_ptr = static_cast(paged_k_cache.data_ptr()); + params.v_ptr = static_cast(paged_v_cache.data_ptr()); + params.o_ptr = static_cast(o.data_ptr()); + params.lse_ptr = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.q_stride_n = q.stride(0); + params.q_stride_h = q.stride(1); + params.o_stride_n = o.stride(0); + params.o_stride_h = o.stride(1); + if (kv_layout == QKVLayout::kNHD) { + // (num_pages, page_size, num_heads, head_dim) + params.k_stride_n = paged_k_cache.stride(1); + params.k_stride_h = paged_k_cache.stride(2); + params.v_stride_n = paged_v_cache.stride(1); + params.v_stride_h = paged_v_cache.stride(2); + } else { + // (num_pages, num_heads, page_size, head_dim) + params.k_stride_h = paged_k_cache.stride(1); + params.k_stride_n = paged_k_cache.stride(2); + params.v_stride_h = paged_v_cache.stride(1); + params.v_stride_n = paged_v_cache.stride(2); + } + params.nnz_qo = q.size(0); + params.num_qo_heads = q.size(1); + params.num_kv_heads = num_kv_heads; + params.group_size = params.num_qo_heads / num_kv_heads; + params.page_size = page_size; + params.window_left = window_left; + params.causal = mask_mode_code == 1; + params.qo_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_tile_indices_offset); + params.qo_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_indptr_offset); + params.kv_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_indptr_offset); + params.qo_lens = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_len_offset); + params.kv_lens = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_len_offset); + params.head_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.head_indices_offset); + params.work_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.work_indptr_offset); + params.batch_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.batch_indices_offset); + params.kv_indices = static_cast(paged_kv_indices.data_ptr()); + + ADDITIONAL_PARAMS_SETTER + + bool same_schedule_for_all_heads = plan_info.same_schedule_for_all_heads; + DISPATCH_BOOL(same_schedule_for_all_heads, SAME_SCHEDULER_FOR_ALL_HEADS, [&] { + cudaError_t status = BatchPrefillWithPagedKVCacheDispatched< + HEAD_DIM_QK, HEAD_DIM_VO, MASK_MODE, USE_SLIDING_WINDOW, SAME_SCHEDULER_FOR_ALL_HEADS, + AttentionVariant>(params, enable_pdl, stream); + TORCH_CHECK(status == cudaSuccess, + "BatchPrefillWithPagedKVCacheSM90Run failed with error: ", + cudaGetErrorString(status)); + return true; + }); + }); +} diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_sm90_config.inc b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_sm90_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..9e64c51ba1efa813f75c40719118a11d5d132ce6 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_sm90_config.inc @@ -0,0 +1,135 @@ +#pragma once +#include +#include +#include +#include +#include +#include "pytorch_extension_utils.h" + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_prefix_len_ptr, std::optional maybe_token_pos_in_items_ptr, std::optional maybe_max_item_len_ptr, double logits_soft_cap, double sm_scale, int64_t token_pos_in_items_len +#define ADDITIONAL_PARAMS_SETTER params.additional_params.maybe_prefix_len_ptr = maybe_prefix_len_ptr ? static_cast(maybe_prefix_len_ptr->data_ptr()): nullptr; \ +params.additional_params.maybe_token_pos_in_items_ptr = maybe_token_pos_in_items_ptr ? static_cast(maybe_token_pos_in_items_ptr->data_ptr()): nullptr; \ +params.additional_params.maybe_max_item_len_ptr = maybe_max_item_len_ptr ? static_cast(maybe_max_item_len_ptr->data_ptr()): nullptr; \ +params.additional_params.logits_soft_cap = logits_soft_cap; \ +params.additional_params.sm_scale = sm_scale; \ +params.additional_params.token_pos_in_items_len = token_pos_in_items_len; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, RaggedParams, PagedParams, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { using AttentionVariant = DefaultAttention; __VA_ARGS__();}) + +using namespace flashinfer; + +using DTypeQ = cutlass_dtype_t; +using DTypeKV = cutlass_dtype_t; +using DTypeO = cutlass_dtype_t; +using IdType = cutlass_dtype_t; + +constexpr int HEAD_DIM_QK = 256; +constexpr int HEAD_DIM_VO = 256; +constexpr auto USE_LOGITS_SOFT_CAP = true; +constexpr auto USE_SLIDING_WINDOW = true; + +struct RaggedParams { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + // The QKV matrices. + DTypeQ* q_ptr; + DTypeKV* k_ptr; + DTypeKV* v_ptr; + DTypeO* o_ptr; + float* lse_ptr; + + IdType* qo_tile_indices; + IdType* qo_indptr; + IdType* kv_indptr; + IdType* qo_lens; + IdType* kv_lens; + IdType* head_indices; + IdType* work_indptr; + IdType* batch_indices; + + struct AdditionalParams { + uint32_t* maybe_prefix_len_ptr; +uint16_t* maybe_token_pos_in_items_ptr; +uint16_t* maybe_max_item_len_ptr; +double logits_soft_cap; +double sm_scale; +int64_t token_pos_in_items_len; + + } additional_params; + + int64_t q_stride_n; + int64_t k_stride_n; + int64_t v_stride_n; + int64_t o_stride_n; + int64_t q_stride_h; + int64_t k_stride_h; + int64_t v_stride_h; + int64_t o_stride_h; + int64_t nnz_qo; + int64_t nnz_kv; + + int head_dim; + int num_qo_heads; + int num_kv_heads; + int group_size; + int window_left; + + bool causal; +}; + +struct PagedParams { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + // The QKV matrices. + DTypeQ* q_ptr; + DTypeKV* k_ptr; + DTypeKV* v_ptr; + DTypeO* o_ptr; + float* lse_ptr; + + IdType* qo_tile_indices; + IdType* qo_indptr; + IdType* kv_indptr; + IdType* kv_indices; + IdType* qo_lens; + IdType* kv_lens; + IdType* head_indices; + IdType* work_indptr; + IdType* batch_indices; + + struct AdditionalParams { + uint32_t* maybe_prefix_len_ptr; +uint16_t* maybe_token_pos_in_items_ptr; +uint16_t* maybe_max_item_len_ptr; +double logits_soft_cap; +double sm_scale; +int64_t token_pos_in_items_len; + + } additional_params; + + int64_t q_stride_n; + int64_t k_stride_n; + int64_t v_stride_n; + int64_t o_stride_n; + int64_t q_stride_h; + int64_t k_stride_h; + int64_t v_stride_h; + int64_t o_stride_h; + int64_t nnz_qo; + + int head_dim; + int num_qo_heads; + int num_kv_heads; + int group_size; + int page_size; + int window_left; + + bool causal; +}; + +#include \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_sm90_jit_pybind.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_sm90_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..b1edd582549e124c776f2cf9b4503f4bf60d16ae --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_sm90_jit_pybind.cu @@ -0,0 +1,46 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "batch_prefill_sm90_config.inc" +#include "pytorch_extension_utils.h" + +at::Tensor BatchPrefillWithKVCacheSM90Plan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, + int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, + int64_t head_dim_vo, bool causal); + +void BatchPrefillWithRaggedKVCacheSM90Run( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +void BatchPrefillWithPagedKVCacheSM90Run( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, + at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Batch-request prefill attention with KV-Cache plan + m.def("plan", BatchPrefillWithKVCacheSM90Plan); + // Batch-request prefill attention with KV-Cache operator + m.def("ragged_run", BatchPrefillWithRaggedKVCacheSM90Run); + // Batch-request prefill attention with KV-Cache operator + m.def("paged_run", BatchPrefillWithPagedKVCacheSM90Run); +} diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill.cu new file mode 100644 index 0000000000000000000000000000000000000000..a51fc7f56a0af7482bd5d2a9d3f672854df69dfd --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill.cu @@ -0,0 +1,330 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include +#include + +#include "batch_prefill_config.inc" +#include "pytorch_conversion_utils.h" +#include "pytorch_extension_utils.h" + +namespace flashinfer { + +template +cudaError_t BatchPrefillWithPagedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v, + float* tmp_s, bool enable_pdl, + cudaStream_t stream); + +template +cudaError_t BatchPrefillWithRaggedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v, + float* tmp_s, bool enable_pdl, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +at::Tensor BatchPrefillWithKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, + int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, + int64_t head_dim_vo, bool causal) { + size_t float_workspace_size_in_bytes = + float_workspace_buffer.size(0) * float_workspace_buffer.element_size(); + size_t int_workspace_size_in_bytes = + int_workspace_buffer.size(0) * int_workspace_buffer.element_size(); + + PrefillPlanInfo plan_info; + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + cudaError_t status = PrefillPlan( + float_workspace_buffer.data_ptr(), float_workspace_size_in_bytes, + int_workspace_buffer.data_ptr(), page_locked_int_workspace_buffer.data_ptr(), + int_workspace_size_in_bytes, plan_info, qo_indptr.data_ptr(), + kv_indptr.data_ptr(), total_num_rows, batch_size, num_qo_heads, num_kv_heads, + head_dim_qk, head_dim_vo, page_size, enable_cuda_graph, /*sizeof_dtype_o=*/2, stream); + + TORCH_CHECK(status == cudaSuccess, + "Failed to plan prefill with error: ", cudaGetErrorString(status)); + + return vec_to_tensor(plan_info.ToVector()); +} + +void BatchPrefillWithRaggedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor k, at::Tensor v, + at::Tensor qo_indptr, at::Tensor kv_indptr, at::Tensor o, + std::optional maybe_lse, int64_t mask_mode_code, + int64_t layout, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + PrefillPlanInfo plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + QKVLayout kv_layout = static_cast(layout); + + int64_t num_qo_heads = q.size(1); + int64_t head_dim_qk = q.size(2); + int64_t num_kv_heads = (kv_layout == QKVLayout::kNHD) ? k.size(1) : k.size(0); + uint32_t q_stride_n = q.stride(0), q_stride_h = q.stride(1), k_stride_n, k_stride_h, v_stride_n, + v_stride_h; + if (kv_layout == QKVLayout::kNHD) { + k_stride_n = k.stride(0); + k_stride_h = k.stride(1); + v_stride_n = v.stride(0); + v_stride_h = v.stride(1); + } else { + k_stride_h = k.stride(0); + k_stride_n = k.stride(1); + v_stride_h = v.stride(0); + v_stride_n = v.stride(1); + } + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == q.size(0), lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == q.size(1), lse.size(1), q.size(1)); + } + + void* float_buffer_ptr = float_workspace_buffer.data_ptr(); + void* int_buffer_ptr = int_workspace_buffer.data_ptr(); + + const MaskMode mask_mode = static_cast(mask_mode_code); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, + RaggedParams, PagedParams, [&] { + RaggedParams params; + + params.q = static_cast(q.data_ptr()); + params.k = static_cast(k.data_ptr()); + params.v = static_cast(v.data_ptr()); + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.q_indptr = static_cast(qo_indptr.data_ptr()); + params.kv_indptr = static_cast(kv_indptr.data_ptr()); + params.num_qo_heads = num_qo_heads; + params.num_kv_heads = num_kv_heads; + params.group_size = uint_fastdiv(num_qo_heads / num_kv_heads); + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.k_stride_n = k_stride_n; + params.k_stride_h = k_stride_h; + params.v_stride_n = v_stride_n; + params.v_stride_h = v_stride_h; + params.window_left = window_left; + + params.request_indices = nullptr; + params.qo_tile_indices = nullptr; + params.kv_tile_indices = nullptr; + params.merge_indptr = nullptr; + params.o_indptr = nullptr; + params.kv_chunk_size_ptr = nullptr; + params.block_valid_mask = nullptr; + params.total_num_rows = nullptr; + params.max_total_num_rows = 0; + params.padded_batch_size = 0; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + DTypeO* tmp_v = nullptr; + float* tmp_s = nullptr; + + params.request_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.request_indices_offset); + params.qo_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_tile_indices_offset); + params.kv_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_tile_indices_offset); + params.o_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.o_indptr_offset); + params.kv_chunk_size_ptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_chunk_size_ptr_offset); + if (plan_info.split_kv) { + params.merge_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.merge_indptr_offset); + tmp_v = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.v_offset); + tmp_s = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.s_offset); + if (plan_info.enable_cuda_graph) { + params.block_valid_mask = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.block_valid_mask_offset); + } + } + params.padded_batch_size = plan_info.padded_batch_size; + params.max_total_num_rows = plan_info.total_num_rows; + if (plan_info.enable_cuda_graph) { + params.total_num_rows = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.total_num_rows_offset); + } + + cudaError_t status = cudaSuccess; + + DISPATCH_CTA_TILE_Q(plan_info.cta_tile_q, CTA_TILE_Q, { + status = flashinfer::BatchPrefillWithRaggedKVCacheDispatched< + CTA_TILE_Q, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + /*use_fp16_qk_reduction=*/USE_FP16_QK_REDUCTION, MASK_MODE, AttentionVariant, + RaggedParams>(params, tmp_v, tmp_s, enable_pdl, stream); + }); + + TORCH_CHECK(status == cudaSuccess, "BatchPrefillWithRaggedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); +} + +void BatchPrefillWithPagedKVCacheRun( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, + at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + PrefillPlanInfo plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + QKVLayout kv_layout = static_cast(layout); + auto device = q.device(); + int64_t batch_size = paged_kv_indptr.size(0) - 1; + int64_t num_qo_heads = q.size(1); + int64_t num_kv_heads, page_size; + uint32_t head_dim_qk = q.size(2); + if (kv_layout == QKVLayout::kHND) { + num_kv_heads = paged_k_cache.size(1); + page_size = paged_k_cache.size(2); + } else { + page_size = paged_k_cache.size(1); + num_kv_heads = paged_k_cache.size(2); + } + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == q.size(0), lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == q.size(1), lse.size(1), q.size(1)); + } + + void* float_buffer_ptr = static_cast(float_workspace_buffer.data_ptr()); + void* int_buffer_ptr = static_cast(int_workspace_buffer.data_ptr()); + + const MaskMode mask_mode = static_cast(mask_mode_code); + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = paged_k_cache.scalar_type(); + + // get q_stride_n and q_stride_h + const auto q_stride_n = q.stride(0); + const auto q_stride_h = q.stride(1); + + // get kv_cache_strides + const int64_t* kv_cache_strides = nullptr; + auto k_strides = paged_k_cache.strides(); + auto v_strides = paged_v_cache.strides(); + TORCH_CHECK(k_strides == v_strides, "k/v strides must be identical"); + kv_cache_strides = k_strides.data(); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, + RaggedParams, PagedParams, [&] { + PagedParams params; + + params.q = static_cast(q.data_ptr()); + paged_kv_t paged_kv( + num_kv_heads, page_size, HEAD_DIM_VO, batch_size, kv_layout, + static_cast(paged_k_cache.data_ptr()), + static_cast(paged_v_cache.data_ptr()), kv_cache_strides, + static_cast(paged_kv_indices.data_ptr()), + static_cast(paged_kv_indptr.data_ptr()), + static_cast(paged_kv_last_page_len.data_ptr())); + params.paged_kv = paged_kv; + params.q_indptr = static_cast(qo_indptr.data_ptr()); + params.o = static_cast(o.data_ptr()); + + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.num_qo_heads = num_qo_heads; + params.group_size = uint_fastdiv(num_qo_heads / paged_kv.num_heads); + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.window_left = window_left; + + params.request_indices = nullptr; + params.qo_tile_indices = nullptr; + params.kv_tile_indices = nullptr; + params.merge_indptr = nullptr; + params.o_indptr = nullptr; + params.kv_chunk_size_ptr = nullptr; + params.block_valid_mask = nullptr; + params.total_num_rows = nullptr; + params.max_total_num_rows = 0; + params.padded_batch_size = 0; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + DTypeO* tmp_v = nullptr; + float* tmp_s = nullptr; + + params.request_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.request_indices_offset); + params.qo_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_tile_indices_offset); + params.kv_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_tile_indices_offset); + params.o_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.o_indptr_offset); + params.kv_chunk_size_ptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_chunk_size_ptr_offset); + if (plan_info.split_kv) { + params.merge_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.merge_indptr_offset); + tmp_v = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.v_offset); + tmp_s = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.s_offset); + if (plan_info.enable_cuda_graph) { + params.block_valid_mask = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.block_valid_mask_offset); + } + } + params.padded_batch_size = plan_info.padded_batch_size; + params.max_total_num_rows = plan_info.total_num_rows; + if (plan_info.enable_cuda_graph) { + params.total_num_rows = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.total_num_rows_offset); + } + + cudaError_t status = cudaSuccess; + + DISPATCH_CTA_TILE_Q(plan_info.cta_tile_q, CTA_TILE_Q, { + status = flashinfer::BatchPrefillWithPagedKVCacheDispatched< + CTA_TILE_Q, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + /*use_fp16_qk_reduction=*/USE_FP16_QK_REDUCTION, MASK_MODE, AttentionVariant, + PagedParams>(params, tmp_v, tmp_s, enable_pdl, stream); + }); + + TORCH_CHECK(status == cudaSuccess, "BatchPrefillWithPagedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); +} diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_config.inc b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..6fe1919986353a75f1299932448ebf2d96ea0baa --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_config.inc @@ -0,0 +1,153 @@ +#pragma once +#include +#include +#include +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_custom_mask, std::optional maybe_mask_indptr, std::optional maybe_alibi_slopes, std::optional maybe_prefix_len_ptr, std::optional maybe_token_pos_in_items_ptr, std::optional maybe_max_item_len_ptr, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta, int64_t token_pos_in_items_len +#define ADDITIONAL_PARAMS_SETTER params.maybe_custom_mask = maybe_custom_mask ? static_cast(maybe_custom_mask->data_ptr()): nullptr; \ +params.maybe_mask_indptr = maybe_mask_indptr ? static_cast(maybe_mask_indptr->data_ptr()): nullptr; \ +params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.maybe_prefix_len_ptr = maybe_prefix_len_ptr ? static_cast(maybe_prefix_len_ptr->data_ptr()): nullptr; \ +params.maybe_token_pos_in_items_ptr = maybe_token_pos_in_items_ptr ? static_cast(maybe_token_pos_in_items_ptr->data_ptr()): nullptr; \ +params.maybe_max_item_len_ptr = maybe_max_item_len_ptr ? static_cast(maybe_max_item_len_ptr->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; \ +params.token_pos_in_items_len = token_pos_in_items_len; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, RaggedParams, PagedParams, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { \ + constexpr auto use_custom_mask = MASK_MODE == MaskMode::kCustom; \ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ + }) + +using namespace flashinfer; + +using DTypeQ = nv_bfloat16; +using DTypeKV = nv_bfloat16; +using DTypeO = nv_bfloat16; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 64; +constexpr int HEAD_DIM_VO = 64; +constexpr bool USE_FP16_QK_REDUCTION = false; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = false; + + +struct RaggedParams { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + DTypeQ* q; + DTypeKV* k; + DTypeKV* v; + IdType* q_indptr; + IdType* kv_indptr; + DTypeO* o; + float* lse; + uint_fastdiv group_size; + + uint8_t* maybe_custom_mask; +int32_t* maybe_mask_indptr; +float* maybe_alibi_slopes; +uint32_t* maybe_prefix_len_ptr; +uint16_t* maybe_token_pos_in_items_ptr; +uint16_t* maybe_max_item_len_ptr; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; +int64_t token_pos_in_items_len; + + uint32_t num_qo_heads; + uint32_t num_kv_heads; + uint32_t q_stride_n; + uint32_t q_stride_h; + uint32_t k_stride_n; + uint32_t k_stride_h; + uint32_t v_stride_n; + uint32_t v_stride_h; + int32_t window_left; + + IdType* request_indices; + IdType* qo_tile_indices; + IdType* kv_tile_indices; + IdType* merge_indptr; + IdType* o_indptr; + IdType* kv_chunk_size_ptr; + bool* block_valid_mask; + uint32_t max_total_num_rows; + uint32_t* total_num_rows; + uint32_t padded_batch_size; + bool partition_kv; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { + return q_indptr[batch_idx + 1] - q_indptr[batch_idx]; + } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return kv_indptr[batch_idx + 1] - kv_indptr[batch_idx]; + } +}; + +struct PagedParams { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + DTypeQ* q; + paged_kv_t paged_kv; + IdType* q_indptr; + DTypeO* o; + float* lse; + uint_fastdiv group_size; + + uint8_t* maybe_custom_mask; +int32_t* maybe_mask_indptr; +float* maybe_alibi_slopes; +uint32_t* maybe_prefix_len_ptr; +uint16_t* maybe_token_pos_in_items_ptr; +uint16_t* maybe_max_item_len_ptr; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; +int64_t token_pos_in_items_len; + + uint32_t num_qo_heads; + IdType q_stride_n; + IdType q_stride_h; + int32_t window_left; + + IdType* request_indices; + IdType* qo_tile_indices; + IdType* kv_tile_indices; + IdType* merge_indptr; + IdType* o_indptr; + bool* block_valid_mask; + IdType* kv_chunk_size_ptr; + uint32_t max_total_num_rows; + uint32_t* total_num_rows; + uint32_t padded_batch_size; + bool partition_kv; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { + return q_indptr[batch_idx + 1] - q_indptr[batch_idx]; + } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return paged_kv.get_length(batch_idx); + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_jit_pybind.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..5421ab1cf910fdfbc867c182c8e22f23e44749ae --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_jit_pybind.cu @@ -0,0 +1,48 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "batch_prefill_config.inc" +#include "pytorch_extension_utils.h" + +at::Tensor BatchPrefillWithKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, + int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, + int64_t head_dim_vo, bool causal); + +void BatchPrefillWithRaggedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor k, at::Tensor v, + at::Tensor qo_indptr, at::Tensor kv_indptr, at::Tensor o, + std::optional maybe_lse, int64_t mask_mode_code, + int64_t layout, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +void BatchPrefillWithPagedKVCacheRun( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, + at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Batch-request prefill attention with KV-Cache plan + m.def("plan", BatchPrefillWithKVCachePlan); + // Batch-request prefill attention with KV-Cache operator + m.def("ragged_run", BatchPrefillWithRaggedKVCacheRun); + // Batch-request prefill attention with KV-Cache operator + m.def("paged_run", BatchPrefillWithPagedKVCacheRun); +} diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_0.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..dc470ac4930d3f74861b25a070ad59708802c8fd --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_0.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kNone == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_1.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..c94a2ee8dc6664dca2e0ffded85071bd6bd61864 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_1.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCausal == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_2.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..cfdc84f24991a8441236dd46f6a83b9dabe31531 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_2.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCustom == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_3.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..ea7d2bb0d4590de5f5d505a6f69f1654f4186ba6 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_3.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kMultiItemScoring == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_0.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..312ae24e057e022000309e4d9dab394e17b24613 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_0.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kNone == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_1.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..9df75df63e5639cedd81990d587d783f2fdc3efb --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_1.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCausal == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_2.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..476f7e93b1d1660b453d5c2d3dc7a4cacf23c24b --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_2.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCustom == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_3.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..ad1770813fab3d00b78cfec767f9df3475d31ea6 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_3.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kMultiItemScoring == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill.cu new file mode 100644 index 0000000000000000000000000000000000000000..a51fc7f56a0af7482bd5d2a9d3f672854df69dfd --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill.cu @@ -0,0 +1,330 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include +#include + +#include "batch_prefill_config.inc" +#include "pytorch_conversion_utils.h" +#include "pytorch_extension_utils.h" + +namespace flashinfer { + +template +cudaError_t BatchPrefillWithPagedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v, + float* tmp_s, bool enable_pdl, + cudaStream_t stream); + +template +cudaError_t BatchPrefillWithRaggedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v, + float* tmp_s, bool enable_pdl, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +at::Tensor BatchPrefillWithKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, + int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, + int64_t head_dim_vo, bool causal) { + size_t float_workspace_size_in_bytes = + float_workspace_buffer.size(0) * float_workspace_buffer.element_size(); + size_t int_workspace_size_in_bytes = + int_workspace_buffer.size(0) * int_workspace_buffer.element_size(); + + PrefillPlanInfo plan_info; + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + cudaError_t status = PrefillPlan( + float_workspace_buffer.data_ptr(), float_workspace_size_in_bytes, + int_workspace_buffer.data_ptr(), page_locked_int_workspace_buffer.data_ptr(), + int_workspace_size_in_bytes, plan_info, qo_indptr.data_ptr(), + kv_indptr.data_ptr(), total_num_rows, batch_size, num_qo_heads, num_kv_heads, + head_dim_qk, head_dim_vo, page_size, enable_cuda_graph, /*sizeof_dtype_o=*/2, stream); + + TORCH_CHECK(status == cudaSuccess, + "Failed to plan prefill with error: ", cudaGetErrorString(status)); + + return vec_to_tensor(plan_info.ToVector()); +} + +void BatchPrefillWithRaggedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor k, at::Tensor v, + at::Tensor qo_indptr, at::Tensor kv_indptr, at::Tensor o, + std::optional maybe_lse, int64_t mask_mode_code, + int64_t layout, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + PrefillPlanInfo plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + QKVLayout kv_layout = static_cast(layout); + + int64_t num_qo_heads = q.size(1); + int64_t head_dim_qk = q.size(2); + int64_t num_kv_heads = (kv_layout == QKVLayout::kNHD) ? k.size(1) : k.size(0); + uint32_t q_stride_n = q.stride(0), q_stride_h = q.stride(1), k_stride_n, k_stride_h, v_stride_n, + v_stride_h; + if (kv_layout == QKVLayout::kNHD) { + k_stride_n = k.stride(0); + k_stride_h = k.stride(1); + v_stride_n = v.stride(0); + v_stride_h = v.stride(1); + } else { + k_stride_h = k.stride(0); + k_stride_n = k.stride(1); + v_stride_h = v.stride(0); + v_stride_n = v.stride(1); + } + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == q.size(0), lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == q.size(1), lse.size(1), q.size(1)); + } + + void* float_buffer_ptr = float_workspace_buffer.data_ptr(); + void* int_buffer_ptr = int_workspace_buffer.data_ptr(); + + const MaskMode mask_mode = static_cast(mask_mode_code); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, + RaggedParams, PagedParams, [&] { + RaggedParams params; + + params.q = static_cast(q.data_ptr()); + params.k = static_cast(k.data_ptr()); + params.v = static_cast(v.data_ptr()); + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.q_indptr = static_cast(qo_indptr.data_ptr()); + params.kv_indptr = static_cast(kv_indptr.data_ptr()); + params.num_qo_heads = num_qo_heads; + params.num_kv_heads = num_kv_heads; + params.group_size = uint_fastdiv(num_qo_heads / num_kv_heads); + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.k_stride_n = k_stride_n; + params.k_stride_h = k_stride_h; + params.v_stride_n = v_stride_n; + params.v_stride_h = v_stride_h; + params.window_left = window_left; + + params.request_indices = nullptr; + params.qo_tile_indices = nullptr; + params.kv_tile_indices = nullptr; + params.merge_indptr = nullptr; + params.o_indptr = nullptr; + params.kv_chunk_size_ptr = nullptr; + params.block_valid_mask = nullptr; + params.total_num_rows = nullptr; + params.max_total_num_rows = 0; + params.padded_batch_size = 0; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + DTypeO* tmp_v = nullptr; + float* tmp_s = nullptr; + + params.request_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.request_indices_offset); + params.qo_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_tile_indices_offset); + params.kv_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_tile_indices_offset); + params.o_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.o_indptr_offset); + params.kv_chunk_size_ptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_chunk_size_ptr_offset); + if (plan_info.split_kv) { + params.merge_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.merge_indptr_offset); + tmp_v = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.v_offset); + tmp_s = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.s_offset); + if (plan_info.enable_cuda_graph) { + params.block_valid_mask = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.block_valid_mask_offset); + } + } + params.padded_batch_size = plan_info.padded_batch_size; + params.max_total_num_rows = plan_info.total_num_rows; + if (plan_info.enable_cuda_graph) { + params.total_num_rows = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.total_num_rows_offset); + } + + cudaError_t status = cudaSuccess; + + DISPATCH_CTA_TILE_Q(plan_info.cta_tile_q, CTA_TILE_Q, { + status = flashinfer::BatchPrefillWithRaggedKVCacheDispatched< + CTA_TILE_Q, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + /*use_fp16_qk_reduction=*/USE_FP16_QK_REDUCTION, MASK_MODE, AttentionVariant, + RaggedParams>(params, tmp_v, tmp_s, enable_pdl, stream); + }); + + TORCH_CHECK(status == cudaSuccess, "BatchPrefillWithRaggedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); +} + +void BatchPrefillWithPagedKVCacheRun( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, + at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + PrefillPlanInfo plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + QKVLayout kv_layout = static_cast(layout); + auto device = q.device(); + int64_t batch_size = paged_kv_indptr.size(0) - 1; + int64_t num_qo_heads = q.size(1); + int64_t num_kv_heads, page_size; + uint32_t head_dim_qk = q.size(2); + if (kv_layout == QKVLayout::kHND) { + num_kv_heads = paged_k_cache.size(1); + page_size = paged_k_cache.size(2); + } else { + page_size = paged_k_cache.size(1); + num_kv_heads = paged_k_cache.size(2); + } + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == q.size(0), lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == q.size(1), lse.size(1), q.size(1)); + } + + void* float_buffer_ptr = static_cast(float_workspace_buffer.data_ptr()); + void* int_buffer_ptr = static_cast(int_workspace_buffer.data_ptr()); + + const MaskMode mask_mode = static_cast(mask_mode_code); + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = paged_k_cache.scalar_type(); + + // get q_stride_n and q_stride_h + const auto q_stride_n = q.stride(0); + const auto q_stride_h = q.stride(1); + + // get kv_cache_strides + const int64_t* kv_cache_strides = nullptr; + auto k_strides = paged_k_cache.strides(); + auto v_strides = paged_v_cache.strides(); + TORCH_CHECK(k_strides == v_strides, "k/v strides must be identical"); + kv_cache_strides = k_strides.data(); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, + RaggedParams, PagedParams, [&] { + PagedParams params; + + params.q = static_cast(q.data_ptr()); + paged_kv_t paged_kv( + num_kv_heads, page_size, HEAD_DIM_VO, batch_size, kv_layout, + static_cast(paged_k_cache.data_ptr()), + static_cast(paged_v_cache.data_ptr()), kv_cache_strides, + static_cast(paged_kv_indices.data_ptr()), + static_cast(paged_kv_indptr.data_ptr()), + static_cast(paged_kv_last_page_len.data_ptr())); + params.paged_kv = paged_kv; + params.q_indptr = static_cast(qo_indptr.data_ptr()); + params.o = static_cast(o.data_ptr()); + + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.num_qo_heads = num_qo_heads; + params.group_size = uint_fastdiv(num_qo_heads / paged_kv.num_heads); + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.window_left = window_left; + + params.request_indices = nullptr; + params.qo_tile_indices = nullptr; + params.kv_tile_indices = nullptr; + params.merge_indptr = nullptr; + params.o_indptr = nullptr; + params.kv_chunk_size_ptr = nullptr; + params.block_valid_mask = nullptr; + params.total_num_rows = nullptr; + params.max_total_num_rows = 0; + params.padded_batch_size = 0; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + DTypeO* tmp_v = nullptr; + float* tmp_s = nullptr; + + params.request_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.request_indices_offset); + params.qo_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_tile_indices_offset); + params.kv_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_tile_indices_offset); + params.o_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.o_indptr_offset); + params.kv_chunk_size_ptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_chunk_size_ptr_offset); + if (plan_info.split_kv) { + params.merge_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.merge_indptr_offset); + tmp_v = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.v_offset); + tmp_s = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.s_offset); + if (plan_info.enable_cuda_graph) { + params.block_valid_mask = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.block_valid_mask_offset); + } + } + params.padded_batch_size = plan_info.padded_batch_size; + params.max_total_num_rows = plan_info.total_num_rows; + if (plan_info.enable_cuda_graph) { + params.total_num_rows = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.total_num_rows_offset); + } + + cudaError_t status = cudaSuccess; + + DISPATCH_CTA_TILE_Q(plan_info.cta_tile_q, CTA_TILE_Q, { + status = flashinfer::BatchPrefillWithPagedKVCacheDispatched< + CTA_TILE_Q, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + /*use_fp16_qk_reduction=*/USE_FP16_QK_REDUCTION, MASK_MODE, AttentionVariant, + PagedParams>(params, tmp_v, tmp_s, enable_pdl, stream); + }); + + TORCH_CHECK(status == cudaSuccess, "BatchPrefillWithPagedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); +} diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_config.inc b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..a72c79168db40f6cd22ba5dff85aaad87db92cb8 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_config.inc @@ -0,0 +1,153 @@ +#pragma once +#include +#include +#include +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_custom_mask, std::optional maybe_mask_indptr, std::optional maybe_alibi_slopes, std::optional maybe_prefix_len_ptr, std::optional maybe_token_pos_in_items_ptr, std::optional maybe_max_item_len_ptr, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta, int64_t token_pos_in_items_len +#define ADDITIONAL_PARAMS_SETTER params.maybe_custom_mask = maybe_custom_mask ? static_cast(maybe_custom_mask->data_ptr()): nullptr; \ +params.maybe_mask_indptr = maybe_mask_indptr ? static_cast(maybe_mask_indptr->data_ptr()): nullptr; \ +params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.maybe_prefix_len_ptr = maybe_prefix_len_ptr ? static_cast(maybe_prefix_len_ptr->data_ptr()): nullptr; \ +params.maybe_token_pos_in_items_ptr = maybe_token_pos_in_items_ptr ? static_cast(maybe_token_pos_in_items_ptr->data_ptr()): nullptr; \ +params.maybe_max_item_len_ptr = maybe_max_item_len_ptr ? static_cast(maybe_max_item_len_ptr->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; \ +params.token_pos_in_items_len = token_pos_in_items_len; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, RaggedParams, PagedParams, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { \ + constexpr auto use_custom_mask = MASK_MODE == MaskMode::kCustom; \ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ + }) + +using namespace flashinfer; + +using DTypeQ = nv_bfloat16; +using DTypeKV = nv_bfloat16; +using DTypeO = nv_bfloat16; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 64; +constexpr int HEAD_DIM_VO = 64; +constexpr bool USE_FP16_QK_REDUCTION = false; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = true; + + +struct RaggedParams { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + DTypeQ* q; + DTypeKV* k; + DTypeKV* v; + IdType* q_indptr; + IdType* kv_indptr; + DTypeO* o; + float* lse; + uint_fastdiv group_size; + + uint8_t* maybe_custom_mask; +int32_t* maybe_mask_indptr; +float* maybe_alibi_slopes; +uint32_t* maybe_prefix_len_ptr; +uint16_t* maybe_token_pos_in_items_ptr; +uint16_t* maybe_max_item_len_ptr; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; +int64_t token_pos_in_items_len; + + uint32_t num_qo_heads; + uint32_t num_kv_heads; + uint32_t q_stride_n; + uint32_t q_stride_h; + uint32_t k_stride_n; + uint32_t k_stride_h; + uint32_t v_stride_n; + uint32_t v_stride_h; + int32_t window_left; + + IdType* request_indices; + IdType* qo_tile_indices; + IdType* kv_tile_indices; + IdType* merge_indptr; + IdType* o_indptr; + IdType* kv_chunk_size_ptr; + bool* block_valid_mask; + uint32_t max_total_num_rows; + uint32_t* total_num_rows; + uint32_t padded_batch_size; + bool partition_kv; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { + return q_indptr[batch_idx + 1] - q_indptr[batch_idx]; + } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return kv_indptr[batch_idx + 1] - kv_indptr[batch_idx]; + } +}; + +struct PagedParams { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + DTypeQ* q; + paged_kv_t paged_kv; + IdType* q_indptr; + DTypeO* o; + float* lse; + uint_fastdiv group_size; + + uint8_t* maybe_custom_mask; +int32_t* maybe_mask_indptr; +float* maybe_alibi_slopes; +uint32_t* maybe_prefix_len_ptr; +uint16_t* maybe_token_pos_in_items_ptr; +uint16_t* maybe_max_item_len_ptr; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; +int64_t token_pos_in_items_len; + + uint32_t num_qo_heads; + IdType q_stride_n; + IdType q_stride_h; + int32_t window_left; + + IdType* request_indices; + IdType* qo_tile_indices; + IdType* kv_tile_indices; + IdType* merge_indptr; + IdType* o_indptr; + bool* block_valid_mask; + IdType* kv_chunk_size_ptr; + uint32_t max_total_num_rows; + uint32_t* total_num_rows; + uint32_t padded_batch_size; + bool partition_kv; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { + return q_indptr[batch_idx + 1] - q_indptr[batch_idx]; + } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return paged_kv.get_length(batch_idx); + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_jit_pybind.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..5421ab1cf910fdfbc867c182c8e22f23e44749ae --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_jit_pybind.cu @@ -0,0 +1,48 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "batch_prefill_config.inc" +#include "pytorch_extension_utils.h" + +at::Tensor BatchPrefillWithKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, + int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, + int64_t head_dim_vo, bool causal); + +void BatchPrefillWithRaggedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor k, at::Tensor v, + at::Tensor qo_indptr, at::Tensor kv_indptr, at::Tensor o, + std::optional maybe_lse, int64_t mask_mode_code, + int64_t layout, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +void BatchPrefillWithPagedKVCacheRun( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, + at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Batch-request prefill attention with KV-Cache plan + m.def("plan", BatchPrefillWithKVCachePlan); + // Batch-request prefill attention with KV-Cache operator + m.def("ragged_run", BatchPrefillWithRaggedKVCacheRun); + // Batch-request prefill attention with KV-Cache operator + m.def("paged_run", BatchPrefillWithPagedKVCacheRun); +} diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_0.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..e18cf3fb9d174fdedc0fd1a8b54cd9bdba181250 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_0.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kNone == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_1.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..2ca8d62b9a6503e0367aa85a39fc5cc164a2e1be --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_1.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCausal == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_2.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..4f16cdb3774a6c0b5727ae57ad2672298a4e1446 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_2.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCustom == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_3.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..5143b75b7cc8d697393aa524693ff750f92bcb82 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_3.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kMultiItemScoring == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_0.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..21713a0ad335444b8200c6481c5b8d4619082752 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_0.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kNone == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_1.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..503fe49906ece459dfe133fc9ee55a1a3ace4625 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_1.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCausal == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_2.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..edcf2c0a482136da1804d7e6d038946b01cb2d49 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_2.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCustom == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_3.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..a2419d83e0324f8b398b949e02a95744b1a64e3b --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_3.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kMultiItemScoring == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill.cu new file mode 100644 index 0000000000000000000000000000000000000000..a51fc7f56a0af7482bd5d2a9d3f672854df69dfd --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill.cu @@ -0,0 +1,330 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include +#include + +#include "batch_prefill_config.inc" +#include "pytorch_conversion_utils.h" +#include "pytorch_extension_utils.h" + +namespace flashinfer { + +template +cudaError_t BatchPrefillWithPagedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v, + float* tmp_s, bool enable_pdl, + cudaStream_t stream); + +template +cudaError_t BatchPrefillWithRaggedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v, + float* tmp_s, bool enable_pdl, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +at::Tensor BatchPrefillWithKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, + int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, + int64_t head_dim_vo, bool causal) { + size_t float_workspace_size_in_bytes = + float_workspace_buffer.size(0) * float_workspace_buffer.element_size(); + size_t int_workspace_size_in_bytes = + int_workspace_buffer.size(0) * int_workspace_buffer.element_size(); + + PrefillPlanInfo plan_info; + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + cudaError_t status = PrefillPlan( + float_workspace_buffer.data_ptr(), float_workspace_size_in_bytes, + int_workspace_buffer.data_ptr(), page_locked_int_workspace_buffer.data_ptr(), + int_workspace_size_in_bytes, plan_info, qo_indptr.data_ptr(), + kv_indptr.data_ptr(), total_num_rows, batch_size, num_qo_heads, num_kv_heads, + head_dim_qk, head_dim_vo, page_size, enable_cuda_graph, /*sizeof_dtype_o=*/2, stream); + + TORCH_CHECK(status == cudaSuccess, + "Failed to plan prefill with error: ", cudaGetErrorString(status)); + + return vec_to_tensor(plan_info.ToVector()); +} + +void BatchPrefillWithRaggedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor k, at::Tensor v, + at::Tensor qo_indptr, at::Tensor kv_indptr, at::Tensor o, + std::optional maybe_lse, int64_t mask_mode_code, + int64_t layout, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + PrefillPlanInfo plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + QKVLayout kv_layout = static_cast(layout); + + int64_t num_qo_heads = q.size(1); + int64_t head_dim_qk = q.size(2); + int64_t num_kv_heads = (kv_layout == QKVLayout::kNHD) ? k.size(1) : k.size(0); + uint32_t q_stride_n = q.stride(0), q_stride_h = q.stride(1), k_stride_n, k_stride_h, v_stride_n, + v_stride_h; + if (kv_layout == QKVLayout::kNHD) { + k_stride_n = k.stride(0); + k_stride_h = k.stride(1); + v_stride_n = v.stride(0); + v_stride_h = v.stride(1); + } else { + k_stride_h = k.stride(0); + k_stride_n = k.stride(1); + v_stride_h = v.stride(0); + v_stride_n = v.stride(1); + } + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == q.size(0), lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == q.size(1), lse.size(1), q.size(1)); + } + + void* float_buffer_ptr = float_workspace_buffer.data_ptr(); + void* int_buffer_ptr = int_workspace_buffer.data_ptr(); + + const MaskMode mask_mode = static_cast(mask_mode_code); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, + RaggedParams, PagedParams, [&] { + RaggedParams params; + + params.q = static_cast(q.data_ptr()); + params.k = static_cast(k.data_ptr()); + params.v = static_cast(v.data_ptr()); + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.q_indptr = static_cast(qo_indptr.data_ptr()); + params.kv_indptr = static_cast(kv_indptr.data_ptr()); + params.num_qo_heads = num_qo_heads; + params.num_kv_heads = num_kv_heads; + params.group_size = uint_fastdiv(num_qo_heads / num_kv_heads); + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.k_stride_n = k_stride_n; + params.k_stride_h = k_stride_h; + params.v_stride_n = v_stride_n; + params.v_stride_h = v_stride_h; + params.window_left = window_left; + + params.request_indices = nullptr; + params.qo_tile_indices = nullptr; + params.kv_tile_indices = nullptr; + params.merge_indptr = nullptr; + params.o_indptr = nullptr; + params.kv_chunk_size_ptr = nullptr; + params.block_valid_mask = nullptr; + params.total_num_rows = nullptr; + params.max_total_num_rows = 0; + params.padded_batch_size = 0; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + DTypeO* tmp_v = nullptr; + float* tmp_s = nullptr; + + params.request_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.request_indices_offset); + params.qo_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_tile_indices_offset); + params.kv_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_tile_indices_offset); + params.o_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.o_indptr_offset); + params.kv_chunk_size_ptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_chunk_size_ptr_offset); + if (plan_info.split_kv) { + params.merge_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.merge_indptr_offset); + tmp_v = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.v_offset); + tmp_s = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.s_offset); + if (plan_info.enable_cuda_graph) { + params.block_valid_mask = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.block_valid_mask_offset); + } + } + params.padded_batch_size = plan_info.padded_batch_size; + params.max_total_num_rows = plan_info.total_num_rows; + if (plan_info.enable_cuda_graph) { + params.total_num_rows = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.total_num_rows_offset); + } + + cudaError_t status = cudaSuccess; + + DISPATCH_CTA_TILE_Q(plan_info.cta_tile_q, CTA_TILE_Q, { + status = flashinfer::BatchPrefillWithRaggedKVCacheDispatched< + CTA_TILE_Q, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + /*use_fp16_qk_reduction=*/USE_FP16_QK_REDUCTION, MASK_MODE, AttentionVariant, + RaggedParams>(params, tmp_v, tmp_s, enable_pdl, stream); + }); + + TORCH_CHECK(status == cudaSuccess, "BatchPrefillWithRaggedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); +} + +void BatchPrefillWithPagedKVCacheRun( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, + at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + PrefillPlanInfo plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + QKVLayout kv_layout = static_cast(layout); + auto device = q.device(); + int64_t batch_size = paged_kv_indptr.size(0) - 1; + int64_t num_qo_heads = q.size(1); + int64_t num_kv_heads, page_size; + uint32_t head_dim_qk = q.size(2); + if (kv_layout == QKVLayout::kHND) { + num_kv_heads = paged_k_cache.size(1); + page_size = paged_k_cache.size(2); + } else { + page_size = paged_k_cache.size(1); + num_kv_heads = paged_k_cache.size(2); + } + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == q.size(0), lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == q.size(1), lse.size(1), q.size(1)); + } + + void* float_buffer_ptr = static_cast(float_workspace_buffer.data_ptr()); + void* int_buffer_ptr = static_cast(int_workspace_buffer.data_ptr()); + + const MaskMode mask_mode = static_cast(mask_mode_code); + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = paged_k_cache.scalar_type(); + + // get q_stride_n and q_stride_h + const auto q_stride_n = q.stride(0); + const auto q_stride_h = q.stride(1); + + // get kv_cache_strides + const int64_t* kv_cache_strides = nullptr; + auto k_strides = paged_k_cache.strides(); + auto v_strides = paged_v_cache.strides(); + TORCH_CHECK(k_strides == v_strides, "k/v strides must be identical"); + kv_cache_strides = k_strides.data(); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, + RaggedParams, PagedParams, [&] { + PagedParams params; + + params.q = static_cast(q.data_ptr()); + paged_kv_t paged_kv( + num_kv_heads, page_size, HEAD_DIM_VO, batch_size, kv_layout, + static_cast(paged_k_cache.data_ptr()), + static_cast(paged_v_cache.data_ptr()), kv_cache_strides, + static_cast(paged_kv_indices.data_ptr()), + static_cast(paged_kv_indptr.data_ptr()), + static_cast(paged_kv_last_page_len.data_ptr())); + params.paged_kv = paged_kv; + params.q_indptr = static_cast(qo_indptr.data_ptr()); + params.o = static_cast(o.data_ptr()); + + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.num_qo_heads = num_qo_heads; + params.group_size = uint_fastdiv(num_qo_heads / paged_kv.num_heads); + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.window_left = window_left; + + params.request_indices = nullptr; + params.qo_tile_indices = nullptr; + params.kv_tile_indices = nullptr; + params.merge_indptr = nullptr; + params.o_indptr = nullptr; + params.kv_chunk_size_ptr = nullptr; + params.block_valid_mask = nullptr; + params.total_num_rows = nullptr; + params.max_total_num_rows = 0; + params.padded_batch_size = 0; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + DTypeO* tmp_v = nullptr; + float* tmp_s = nullptr; + + params.request_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.request_indices_offset); + params.qo_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_tile_indices_offset); + params.kv_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_tile_indices_offset); + params.o_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.o_indptr_offset); + params.kv_chunk_size_ptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_chunk_size_ptr_offset); + if (plan_info.split_kv) { + params.merge_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.merge_indptr_offset); + tmp_v = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.v_offset); + tmp_s = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.s_offset); + if (plan_info.enable_cuda_graph) { + params.block_valid_mask = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.block_valid_mask_offset); + } + } + params.padded_batch_size = plan_info.padded_batch_size; + params.max_total_num_rows = plan_info.total_num_rows; + if (plan_info.enable_cuda_graph) { + params.total_num_rows = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.total_num_rows_offset); + } + + cudaError_t status = cudaSuccess; + + DISPATCH_CTA_TILE_Q(plan_info.cta_tile_q, CTA_TILE_Q, { + status = flashinfer::BatchPrefillWithPagedKVCacheDispatched< + CTA_TILE_Q, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + /*use_fp16_qk_reduction=*/USE_FP16_QK_REDUCTION, MASK_MODE, AttentionVariant, + PagedParams>(params, tmp_v, tmp_s, enable_pdl, stream); + }); + + TORCH_CHECK(status == cudaSuccess, "BatchPrefillWithPagedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); +} diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_config.inc b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..b4199e83e6d37e67555105b8f4846d221a9bab85 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_config.inc @@ -0,0 +1,153 @@ +#pragma once +#include +#include +#include +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_custom_mask, std::optional maybe_mask_indptr, std::optional maybe_alibi_slopes, std::optional maybe_prefix_len_ptr, std::optional maybe_token_pos_in_items_ptr, std::optional maybe_max_item_len_ptr, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta, int64_t token_pos_in_items_len +#define ADDITIONAL_PARAMS_SETTER params.maybe_custom_mask = maybe_custom_mask ? static_cast(maybe_custom_mask->data_ptr()): nullptr; \ +params.maybe_mask_indptr = maybe_mask_indptr ? static_cast(maybe_mask_indptr->data_ptr()): nullptr; \ +params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.maybe_prefix_len_ptr = maybe_prefix_len_ptr ? static_cast(maybe_prefix_len_ptr->data_ptr()): nullptr; \ +params.maybe_token_pos_in_items_ptr = maybe_token_pos_in_items_ptr ? static_cast(maybe_token_pos_in_items_ptr->data_ptr()): nullptr; \ +params.maybe_max_item_len_ptr = maybe_max_item_len_ptr ? static_cast(maybe_max_item_len_ptr->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; \ +params.token_pos_in_items_len = token_pos_in_items_len; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, RaggedParams, PagedParams, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { \ + constexpr auto use_custom_mask = MASK_MODE == MaskMode::kCustom; \ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ + }) + +using namespace flashinfer; + +using DTypeQ = nv_bfloat16; +using DTypeKV = __nv_fp8_e4m3; +using DTypeO = nv_bfloat16; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 128; +constexpr int HEAD_DIM_VO = 128; +constexpr bool USE_FP16_QK_REDUCTION = false; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = false; + + +struct RaggedParams { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + DTypeQ* q; + DTypeKV* k; + DTypeKV* v; + IdType* q_indptr; + IdType* kv_indptr; + DTypeO* o; + float* lse; + uint_fastdiv group_size; + + uint8_t* maybe_custom_mask; +int32_t* maybe_mask_indptr; +float* maybe_alibi_slopes; +uint32_t* maybe_prefix_len_ptr; +uint16_t* maybe_token_pos_in_items_ptr; +uint16_t* maybe_max_item_len_ptr; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; +int64_t token_pos_in_items_len; + + uint32_t num_qo_heads; + uint32_t num_kv_heads; + uint32_t q_stride_n; + uint32_t q_stride_h; + uint32_t k_stride_n; + uint32_t k_stride_h; + uint32_t v_stride_n; + uint32_t v_stride_h; + int32_t window_left; + + IdType* request_indices; + IdType* qo_tile_indices; + IdType* kv_tile_indices; + IdType* merge_indptr; + IdType* o_indptr; + IdType* kv_chunk_size_ptr; + bool* block_valid_mask; + uint32_t max_total_num_rows; + uint32_t* total_num_rows; + uint32_t padded_batch_size; + bool partition_kv; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { + return q_indptr[batch_idx + 1] - q_indptr[batch_idx]; + } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return kv_indptr[batch_idx + 1] - kv_indptr[batch_idx]; + } +}; + +struct PagedParams { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + DTypeQ* q; + paged_kv_t paged_kv; + IdType* q_indptr; + DTypeO* o; + float* lse; + uint_fastdiv group_size; + + uint8_t* maybe_custom_mask; +int32_t* maybe_mask_indptr; +float* maybe_alibi_slopes; +uint32_t* maybe_prefix_len_ptr; +uint16_t* maybe_token_pos_in_items_ptr; +uint16_t* maybe_max_item_len_ptr; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; +int64_t token_pos_in_items_len; + + uint32_t num_qo_heads; + IdType q_stride_n; + IdType q_stride_h; + int32_t window_left; + + IdType* request_indices; + IdType* qo_tile_indices; + IdType* kv_tile_indices; + IdType* merge_indptr; + IdType* o_indptr; + bool* block_valid_mask; + IdType* kv_chunk_size_ptr; + uint32_t max_total_num_rows; + uint32_t* total_num_rows; + uint32_t padded_batch_size; + bool partition_kv; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { + return q_indptr[batch_idx + 1] - q_indptr[batch_idx]; + } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return paged_kv.get_length(batch_idx); + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_jit_pybind.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..5421ab1cf910fdfbc867c182c8e22f23e44749ae --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_jit_pybind.cu @@ -0,0 +1,48 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "batch_prefill_config.inc" +#include "pytorch_extension_utils.h" + +at::Tensor BatchPrefillWithKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, + int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, + int64_t head_dim_vo, bool causal); + +void BatchPrefillWithRaggedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor k, at::Tensor v, + at::Tensor qo_indptr, at::Tensor kv_indptr, at::Tensor o, + std::optional maybe_lse, int64_t mask_mode_code, + int64_t layout, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +void BatchPrefillWithPagedKVCacheRun( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, + at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Batch-request prefill attention with KV-Cache plan + m.def("plan", BatchPrefillWithKVCachePlan); + // Batch-request prefill attention with KV-Cache operator + m.def("ragged_run", BatchPrefillWithRaggedKVCacheRun); + // Batch-request prefill attention with KV-Cache operator + m.def("paged_run", BatchPrefillWithPagedKVCacheRun); +} diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_0.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..4c306789cc8cbc7b47b4533007bd6ad51ab9a9e1 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_0.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kNone == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 128, 128, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 128, 128, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 128, 128, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_1.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..d40fe782b6dd930e6e8548d9c15df1dacaee3b8b --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_1.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCausal == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 128, 128, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 128, 128, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 128, 128, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_2.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..16e8d67b528952a4650148a9e7119057667f18f6 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_2.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCustom == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 128, 128, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 128, 128, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 128, 128, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_3.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..7faebee71230224faa5fba4ec9709f8d5f811ca2 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_3.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kMultiItemScoring == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 128, 128, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 128, 128, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 128, 128, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_0.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..38d25465ee0a5e231720942d0712969244567fbe --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_0.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kNone == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 128, 128, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 128, 128, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 128, 128, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_1.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..f18906355ac04ee88e405d998e5d3aa211884b1c --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_1.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCausal == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 128, 128, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 128, 128, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 128, 128, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_2.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..bf34ea033aefab899927a75acc5ed1a6e7a7b310 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_2.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCustom == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 128, 128, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 128, 128, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 128, 128, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_3.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..6b6e36ad0a7464e476520a6e6761370b217dfcbb --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_3.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kMultiItemScoring == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 128, 128, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 128, 128, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 128, 128, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill.cu new file mode 100644 index 0000000000000000000000000000000000000000..a51fc7f56a0af7482bd5d2a9d3f672854df69dfd --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill.cu @@ -0,0 +1,330 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include +#include + +#include "batch_prefill_config.inc" +#include "pytorch_conversion_utils.h" +#include "pytorch_extension_utils.h" + +namespace flashinfer { + +template +cudaError_t BatchPrefillWithPagedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v, + float* tmp_s, bool enable_pdl, + cudaStream_t stream); + +template +cudaError_t BatchPrefillWithRaggedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v, + float* tmp_s, bool enable_pdl, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +at::Tensor BatchPrefillWithKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, + int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, + int64_t head_dim_vo, bool causal) { + size_t float_workspace_size_in_bytes = + float_workspace_buffer.size(0) * float_workspace_buffer.element_size(); + size_t int_workspace_size_in_bytes = + int_workspace_buffer.size(0) * int_workspace_buffer.element_size(); + + PrefillPlanInfo plan_info; + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + cudaError_t status = PrefillPlan( + float_workspace_buffer.data_ptr(), float_workspace_size_in_bytes, + int_workspace_buffer.data_ptr(), page_locked_int_workspace_buffer.data_ptr(), + int_workspace_size_in_bytes, plan_info, qo_indptr.data_ptr(), + kv_indptr.data_ptr(), total_num_rows, batch_size, num_qo_heads, num_kv_heads, + head_dim_qk, head_dim_vo, page_size, enable_cuda_graph, /*sizeof_dtype_o=*/2, stream); + + TORCH_CHECK(status == cudaSuccess, + "Failed to plan prefill with error: ", cudaGetErrorString(status)); + + return vec_to_tensor(plan_info.ToVector()); +} + +void BatchPrefillWithRaggedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor k, at::Tensor v, + at::Tensor qo_indptr, at::Tensor kv_indptr, at::Tensor o, + std::optional maybe_lse, int64_t mask_mode_code, + int64_t layout, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + PrefillPlanInfo plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + QKVLayout kv_layout = static_cast(layout); + + int64_t num_qo_heads = q.size(1); + int64_t head_dim_qk = q.size(2); + int64_t num_kv_heads = (kv_layout == QKVLayout::kNHD) ? k.size(1) : k.size(0); + uint32_t q_stride_n = q.stride(0), q_stride_h = q.stride(1), k_stride_n, k_stride_h, v_stride_n, + v_stride_h; + if (kv_layout == QKVLayout::kNHD) { + k_stride_n = k.stride(0); + k_stride_h = k.stride(1); + v_stride_n = v.stride(0); + v_stride_h = v.stride(1); + } else { + k_stride_h = k.stride(0); + k_stride_n = k.stride(1); + v_stride_h = v.stride(0); + v_stride_n = v.stride(1); + } + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == q.size(0), lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == q.size(1), lse.size(1), q.size(1)); + } + + void* float_buffer_ptr = float_workspace_buffer.data_ptr(); + void* int_buffer_ptr = int_workspace_buffer.data_ptr(); + + const MaskMode mask_mode = static_cast(mask_mode_code); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, + RaggedParams, PagedParams, [&] { + RaggedParams params; + + params.q = static_cast(q.data_ptr()); + params.k = static_cast(k.data_ptr()); + params.v = static_cast(v.data_ptr()); + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.q_indptr = static_cast(qo_indptr.data_ptr()); + params.kv_indptr = static_cast(kv_indptr.data_ptr()); + params.num_qo_heads = num_qo_heads; + params.num_kv_heads = num_kv_heads; + params.group_size = uint_fastdiv(num_qo_heads / num_kv_heads); + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.k_stride_n = k_stride_n; + params.k_stride_h = k_stride_h; + params.v_stride_n = v_stride_n; + params.v_stride_h = v_stride_h; + params.window_left = window_left; + + params.request_indices = nullptr; + params.qo_tile_indices = nullptr; + params.kv_tile_indices = nullptr; + params.merge_indptr = nullptr; + params.o_indptr = nullptr; + params.kv_chunk_size_ptr = nullptr; + params.block_valid_mask = nullptr; + params.total_num_rows = nullptr; + params.max_total_num_rows = 0; + params.padded_batch_size = 0; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + DTypeO* tmp_v = nullptr; + float* tmp_s = nullptr; + + params.request_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.request_indices_offset); + params.qo_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_tile_indices_offset); + params.kv_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_tile_indices_offset); + params.o_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.o_indptr_offset); + params.kv_chunk_size_ptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_chunk_size_ptr_offset); + if (plan_info.split_kv) { + params.merge_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.merge_indptr_offset); + tmp_v = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.v_offset); + tmp_s = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.s_offset); + if (plan_info.enable_cuda_graph) { + params.block_valid_mask = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.block_valid_mask_offset); + } + } + params.padded_batch_size = plan_info.padded_batch_size; + params.max_total_num_rows = plan_info.total_num_rows; + if (plan_info.enable_cuda_graph) { + params.total_num_rows = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.total_num_rows_offset); + } + + cudaError_t status = cudaSuccess; + + DISPATCH_CTA_TILE_Q(plan_info.cta_tile_q, CTA_TILE_Q, { + status = flashinfer::BatchPrefillWithRaggedKVCacheDispatched< + CTA_TILE_Q, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + /*use_fp16_qk_reduction=*/USE_FP16_QK_REDUCTION, MASK_MODE, AttentionVariant, + RaggedParams>(params, tmp_v, tmp_s, enable_pdl, stream); + }); + + TORCH_CHECK(status == cudaSuccess, "BatchPrefillWithRaggedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); +} + +void BatchPrefillWithPagedKVCacheRun( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, + at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + PrefillPlanInfo plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + QKVLayout kv_layout = static_cast(layout); + auto device = q.device(); + int64_t batch_size = paged_kv_indptr.size(0) - 1; + int64_t num_qo_heads = q.size(1); + int64_t num_kv_heads, page_size; + uint32_t head_dim_qk = q.size(2); + if (kv_layout == QKVLayout::kHND) { + num_kv_heads = paged_k_cache.size(1); + page_size = paged_k_cache.size(2); + } else { + page_size = paged_k_cache.size(1); + num_kv_heads = paged_k_cache.size(2); + } + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == q.size(0), lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == q.size(1), lse.size(1), q.size(1)); + } + + void* float_buffer_ptr = static_cast(float_workspace_buffer.data_ptr()); + void* int_buffer_ptr = static_cast(int_workspace_buffer.data_ptr()); + + const MaskMode mask_mode = static_cast(mask_mode_code); + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = paged_k_cache.scalar_type(); + + // get q_stride_n and q_stride_h + const auto q_stride_n = q.stride(0); + const auto q_stride_h = q.stride(1); + + // get kv_cache_strides + const int64_t* kv_cache_strides = nullptr; + auto k_strides = paged_k_cache.strides(); + auto v_strides = paged_v_cache.strides(); + TORCH_CHECK(k_strides == v_strides, "k/v strides must be identical"); + kv_cache_strides = k_strides.data(); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, + RaggedParams, PagedParams, [&] { + PagedParams params; + + params.q = static_cast(q.data_ptr()); + paged_kv_t paged_kv( + num_kv_heads, page_size, HEAD_DIM_VO, batch_size, kv_layout, + static_cast(paged_k_cache.data_ptr()), + static_cast(paged_v_cache.data_ptr()), kv_cache_strides, + static_cast(paged_kv_indices.data_ptr()), + static_cast(paged_kv_indptr.data_ptr()), + static_cast(paged_kv_last_page_len.data_ptr())); + params.paged_kv = paged_kv; + params.q_indptr = static_cast(qo_indptr.data_ptr()); + params.o = static_cast(o.data_ptr()); + + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.num_qo_heads = num_qo_heads; + params.group_size = uint_fastdiv(num_qo_heads / paged_kv.num_heads); + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.window_left = window_left; + + params.request_indices = nullptr; + params.qo_tile_indices = nullptr; + params.kv_tile_indices = nullptr; + params.merge_indptr = nullptr; + params.o_indptr = nullptr; + params.kv_chunk_size_ptr = nullptr; + params.block_valid_mask = nullptr; + params.total_num_rows = nullptr; + params.max_total_num_rows = 0; + params.padded_batch_size = 0; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + DTypeO* tmp_v = nullptr; + float* tmp_s = nullptr; + + params.request_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.request_indices_offset); + params.qo_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_tile_indices_offset); + params.kv_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_tile_indices_offset); + params.o_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.o_indptr_offset); + params.kv_chunk_size_ptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_chunk_size_ptr_offset); + if (plan_info.split_kv) { + params.merge_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.merge_indptr_offset); + tmp_v = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.v_offset); + tmp_s = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.s_offset); + if (plan_info.enable_cuda_graph) { + params.block_valid_mask = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.block_valid_mask_offset); + } + } + params.padded_batch_size = plan_info.padded_batch_size; + params.max_total_num_rows = plan_info.total_num_rows; + if (plan_info.enable_cuda_graph) { + params.total_num_rows = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.total_num_rows_offset); + } + + cudaError_t status = cudaSuccess; + + DISPATCH_CTA_TILE_Q(plan_info.cta_tile_q, CTA_TILE_Q, { + status = flashinfer::BatchPrefillWithPagedKVCacheDispatched< + CTA_TILE_Q, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + /*use_fp16_qk_reduction=*/USE_FP16_QK_REDUCTION, MASK_MODE, AttentionVariant, + PagedParams>(params, tmp_v, tmp_s, enable_pdl, stream); + }); + + TORCH_CHECK(status == cudaSuccess, "BatchPrefillWithPagedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); +} diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_config.inc b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..04f960d473061a968a5a9a1fb6189db76b671349 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_config.inc @@ -0,0 +1,153 @@ +#pragma once +#include +#include +#include +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_custom_mask, std::optional maybe_mask_indptr, std::optional maybe_alibi_slopes, std::optional maybe_prefix_len_ptr, std::optional maybe_token_pos_in_items_ptr, std::optional maybe_max_item_len_ptr, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta, int64_t token_pos_in_items_len +#define ADDITIONAL_PARAMS_SETTER params.maybe_custom_mask = maybe_custom_mask ? static_cast(maybe_custom_mask->data_ptr()): nullptr; \ +params.maybe_mask_indptr = maybe_mask_indptr ? static_cast(maybe_mask_indptr->data_ptr()): nullptr; \ +params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.maybe_prefix_len_ptr = maybe_prefix_len_ptr ? static_cast(maybe_prefix_len_ptr->data_ptr()): nullptr; \ +params.maybe_token_pos_in_items_ptr = maybe_token_pos_in_items_ptr ? static_cast(maybe_token_pos_in_items_ptr->data_ptr()): nullptr; \ +params.maybe_max_item_len_ptr = maybe_max_item_len_ptr ? static_cast(maybe_max_item_len_ptr->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; \ +params.token_pos_in_items_len = token_pos_in_items_len; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, RaggedParams, PagedParams, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { \ + constexpr auto use_custom_mask = MASK_MODE == MaskMode::kCustom; \ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ + }) + +using namespace flashinfer; + +using DTypeQ = nv_bfloat16; +using DTypeKV = __nv_fp8_e4m3; +using DTypeO = nv_bfloat16; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 256; +constexpr int HEAD_DIM_VO = 256; +constexpr bool USE_FP16_QK_REDUCTION = false; +constexpr auto USE_LOGITS_SOFT_CAP = true; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = true; + + +struct RaggedParams { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + DTypeQ* q; + DTypeKV* k; + DTypeKV* v; + IdType* q_indptr; + IdType* kv_indptr; + DTypeO* o; + float* lse; + uint_fastdiv group_size; + + uint8_t* maybe_custom_mask; +int32_t* maybe_mask_indptr; +float* maybe_alibi_slopes; +uint32_t* maybe_prefix_len_ptr; +uint16_t* maybe_token_pos_in_items_ptr; +uint16_t* maybe_max_item_len_ptr; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; +int64_t token_pos_in_items_len; + + uint32_t num_qo_heads; + uint32_t num_kv_heads; + uint32_t q_stride_n; + uint32_t q_stride_h; + uint32_t k_stride_n; + uint32_t k_stride_h; + uint32_t v_stride_n; + uint32_t v_stride_h; + int32_t window_left; + + IdType* request_indices; + IdType* qo_tile_indices; + IdType* kv_tile_indices; + IdType* merge_indptr; + IdType* o_indptr; + IdType* kv_chunk_size_ptr; + bool* block_valid_mask; + uint32_t max_total_num_rows; + uint32_t* total_num_rows; + uint32_t padded_batch_size; + bool partition_kv; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { + return q_indptr[batch_idx + 1] - q_indptr[batch_idx]; + } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return kv_indptr[batch_idx + 1] - kv_indptr[batch_idx]; + } +}; + +struct PagedParams { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + DTypeQ* q; + paged_kv_t paged_kv; + IdType* q_indptr; + DTypeO* o; + float* lse; + uint_fastdiv group_size; + + uint8_t* maybe_custom_mask; +int32_t* maybe_mask_indptr; +float* maybe_alibi_slopes; +uint32_t* maybe_prefix_len_ptr; +uint16_t* maybe_token_pos_in_items_ptr; +uint16_t* maybe_max_item_len_ptr; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; +int64_t token_pos_in_items_len; + + uint32_t num_qo_heads; + IdType q_stride_n; + IdType q_stride_h; + int32_t window_left; + + IdType* request_indices; + IdType* qo_tile_indices; + IdType* kv_tile_indices; + IdType* merge_indptr; + IdType* o_indptr; + bool* block_valid_mask; + IdType* kv_chunk_size_ptr; + uint32_t max_total_num_rows; + uint32_t* total_num_rows; + uint32_t padded_batch_size; + bool partition_kv; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { + return q_indptr[batch_idx + 1] - q_indptr[batch_idx]; + } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return paged_kv.get_length(batch_idx); + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_jit_pybind.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..5421ab1cf910fdfbc867c182c8e22f23e44749ae --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_jit_pybind.cu @@ -0,0 +1,48 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "batch_prefill_config.inc" +#include "pytorch_extension_utils.h" + +at::Tensor BatchPrefillWithKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, + int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, + int64_t head_dim_vo, bool causal); + +void BatchPrefillWithRaggedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor k, at::Tensor v, + at::Tensor qo_indptr, at::Tensor kv_indptr, at::Tensor o, + std::optional maybe_lse, int64_t mask_mode_code, + int64_t layout, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +void BatchPrefillWithPagedKVCacheRun( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, + at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Batch-request prefill attention with KV-Cache plan + m.def("plan", BatchPrefillWithKVCachePlan); + // Batch-request prefill attention with KV-Cache operator + m.def("ragged_run", BatchPrefillWithRaggedKVCacheRun); + // Batch-request prefill attention with KV-Cache operator + m.def("paged_run", BatchPrefillWithPagedKVCacheRun); +} diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_paged_kernel_mask_0.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_paged_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..264348261256ffc911b50c9fd7237441faec122f --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_paged_kernel_mask_0.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kNone == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 256, 256, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 256, 256, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 256, 256, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_paged_kernel_mask_1.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_paged_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..9937900da6dca1744966105d9de12f7321ae00e4 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_paged_kernel_mask_1.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCausal == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 256, 256, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 256, 256, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 256, 256, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_paged_kernel_mask_2.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_paged_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..f61a91898a844ffbce56d04112d87a9e078243d8 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_paged_kernel_mask_2.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCustom == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 256, 256, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 256, 256, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 256, 256, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_paged_kernel_mask_3.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_paged_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..e258a720abb5b0a909466b307c58f19ddd43c157 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_paged_kernel_mask_3.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kMultiItemScoring == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 256, 256, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 256, 256, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 256, 256, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_ragged_kernel_mask_0.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_ragged_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..ff099d149943dfad67e18d6b600d3f5781bf5d02 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_ragged_kernel_mask_0.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kNone == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 256, 256, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 256, 256, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 256, 256, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_ragged_kernel_mask_1.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_ragged_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..838dcf44be11d68bb2c48ddbc70c8b42c62d4f93 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_ragged_kernel_mask_1.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCausal == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 256, 256, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 256, 256, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 256, 256, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_ragged_kernel_mask_2.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_ragged_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..9735995235347e9cf67269a5d532eaf67df02b14 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_ragged_kernel_mask_2.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCustom == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 256, 256, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 256, 256, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 256, 256, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_ragged_kernel_mask_3.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_ragged_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..255af5bee1786391bc5be516b8f8c9c051133a41 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_ragged_kernel_mask_3.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kMultiItemScoring == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 256, 256, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 256, 256, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 256, 256, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill.cu new file mode 100644 index 0000000000000000000000000000000000000000..a51fc7f56a0af7482bd5d2a9d3f672854df69dfd --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill.cu @@ -0,0 +1,330 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include +#include + +#include "batch_prefill_config.inc" +#include "pytorch_conversion_utils.h" +#include "pytorch_extension_utils.h" + +namespace flashinfer { + +template +cudaError_t BatchPrefillWithPagedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v, + float* tmp_s, bool enable_pdl, + cudaStream_t stream); + +template +cudaError_t BatchPrefillWithRaggedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v, + float* tmp_s, bool enable_pdl, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +at::Tensor BatchPrefillWithKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, + int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, + int64_t head_dim_vo, bool causal) { + size_t float_workspace_size_in_bytes = + float_workspace_buffer.size(0) * float_workspace_buffer.element_size(); + size_t int_workspace_size_in_bytes = + int_workspace_buffer.size(0) * int_workspace_buffer.element_size(); + + PrefillPlanInfo plan_info; + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + cudaError_t status = PrefillPlan( + float_workspace_buffer.data_ptr(), float_workspace_size_in_bytes, + int_workspace_buffer.data_ptr(), page_locked_int_workspace_buffer.data_ptr(), + int_workspace_size_in_bytes, plan_info, qo_indptr.data_ptr(), + kv_indptr.data_ptr(), total_num_rows, batch_size, num_qo_heads, num_kv_heads, + head_dim_qk, head_dim_vo, page_size, enable_cuda_graph, /*sizeof_dtype_o=*/2, stream); + + TORCH_CHECK(status == cudaSuccess, + "Failed to plan prefill with error: ", cudaGetErrorString(status)); + + return vec_to_tensor(plan_info.ToVector()); +} + +void BatchPrefillWithRaggedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor k, at::Tensor v, + at::Tensor qo_indptr, at::Tensor kv_indptr, at::Tensor o, + std::optional maybe_lse, int64_t mask_mode_code, + int64_t layout, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + PrefillPlanInfo plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + QKVLayout kv_layout = static_cast(layout); + + int64_t num_qo_heads = q.size(1); + int64_t head_dim_qk = q.size(2); + int64_t num_kv_heads = (kv_layout == QKVLayout::kNHD) ? k.size(1) : k.size(0); + uint32_t q_stride_n = q.stride(0), q_stride_h = q.stride(1), k_stride_n, k_stride_h, v_stride_n, + v_stride_h; + if (kv_layout == QKVLayout::kNHD) { + k_stride_n = k.stride(0); + k_stride_h = k.stride(1); + v_stride_n = v.stride(0); + v_stride_h = v.stride(1); + } else { + k_stride_h = k.stride(0); + k_stride_n = k.stride(1); + v_stride_h = v.stride(0); + v_stride_n = v.stride(1); + } + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == q.size(0), lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == q.size(1), lse.size(1), q.size(1)); + } + + void* float_buffer_ptr = float_workspace_buffer.data_ptr(); + void* int_buffer_ptr = int_workspace_buffer.data_ptr(); + + const MaskMode mask_mode = static_cast(mask_mode_code); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, + RaggedParams, PagedParams, [&] { + RaggedParams params; + + params.q = static_cast(q.data_ptr()); + params.k = static_cast(k.data_ptr()); + params.v = static_cast(v.data_ptr()); + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.q_indptr = static_cast(qo_indptr.data_ptr()); + params.kv_indptr = static_cast(kv_indptr.data_ptr()); + params.num_qo_heads = num_qo_heads; + params.num_kv_heads = num_kv_heads; + params.group_size = uint_fastdiv(num_qo_heads / num_kv_heads); + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.k_stride_n = k_stride_n; + params.k_stride_h = k_stride_h; + params.v_stride_n = v_stride_n; + params.v_stride_h = v_stride_h; + params.window_left = window_left; + + params.request_indices = nullptr; + params.qo_tile_indices = nullptr; + params.kv_tile_indices = nullptr; + params.merge_indptr = nullptr; + params.o_indptr = nullptr; + params.kv_chunk_size_ptr = nullptr; + params.block_valid_mask = nullptr; + params.total_num_rows = nullptr; + params.max_total_num_rows = 0; + params.padded_batch_size = 0; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + DTypeO* tmp_v = nullptr; + float* tmp_s = nullptr; + + params.request_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.request_indices_offset); + params.qo_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_tile_indices_offset); + params.kv_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_tile_indices_offset); + params.o_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.o_indptr_offset); + params.kv_chunk_size_ptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_chunk_size_ptr_offset); + if (plan_info.split_kv) { + params.merge_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.merge_indptr_offset); + tmp_v = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.v_offset); + tmp_s = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.s_offset); + if (plan_info.enable_cuda_graph) { + params.block_valid_mask = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.block_valid_mask_offset); + } + } + params.padded_batch_size = plan_info.padded_batch_size; + params.max_total_num_rows = plan_info.total_num_rows; + if (plan_info.enable_cuda_graph) { + params.total_num_rows = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.total_num_rows_offset); + } + + cudaError_t status = cudaSuccess; + + DISPATCH_CTA_TILE_Q(plan_info.cta_tile_q, CTA_TILE_Q, { + status = flashinfer::BatchPrefillWithRaggedKVCacheDispatched< + CTA_TILE_Q, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + /*use_fp16_qk_reduction=*/USE_FP16_QK_REDUCTION, MASK_MODE, AttentionVariant, + RaggedParams>(params, tmp_v, tmp_s, enable_pdl, stream); + }); + + TORCH_CHECK(status == cudaSuccess, "BatchPrefillWithRaggedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); +} + +void BatchPrefillWithPagedKVCacheRun( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, + at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + PrefillPlanInfo plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + QKVLayout kv_layout = static_cast(layout); + auto device = q.device(); + int64_t batch_size = paged_kv_indptr.size(0) - 1; + int64_t num_qo_heads = q.size(1); + int64_t num_kv_heads, page_size; + uint32_t head_dim_qk = q.size(2); + if (kv_layout == QKVLayout::kHND) { + num_kv_heads = paged_k_cache.size(1); + page_size = paged_k_cache.size(2); + } else { + page_size = paged_k_cache.size(1); + num_kv_heads = paged_k_cache.size(2); + } + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == q.size(0), lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == q.size(1), lse.size(1), q.size(1)); + } + + void* float_buffer_ptr = static_cast(float_workspace_buffer.data_ptr()); + void* int_buffer_ptr = static_cast(int_workspace_buffer.data_ptr()); + + const MaskMode mask_mode = static_cast(mask_mode_code); + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = paged_k_cache.scalar_type(); + + // get q_stride_n and q_stride_h + const auto q_stride_n = q.stride(0); + const auto q_stride_h = q.stride(1); + + // get kv_cache_strides + const int64_t* kv_cache_strides = nullptr; + auto k_strides = paged_k_cache.strides(); + auto v_strides = paged_v_cache.strides(); + TORCH_CHECK(k_strides == v_strides, "k/v strides must be identical"); + kv_cache_strides = k_strides.data(); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, + RaggedParams, PagedParams, [&] { + PagedParams params; + + params.q = static_cast(q.data_ptr()); + paged_kv_t paged_kv( + num_kv_heads, page_size, HEAD_DIM_VO, batch_size, kv_layout, + static_cast(paged_k_cache.data_ptr()), + static_cast(paged_v_cache.data_ptr()), kv_cache_strides, + static_cast(paged_kv_indices.data_ptr()), + static_cast(paged_kv_indptr.data_ptr()), + static_cast(paged_kv_last_page_len.data_ptr())); + params.paged_kv = paged_kv; + params.q_indptr = static_cast(qo_indptr.data_ptr()); + params.o = static_cast(o.data_ptr()); + + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.num_qo_heads = num_qo_heads; + params.group_size = uint_fastdiv(num_qo_heads / paged_kv.num_heads); + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.window_left = window_left; + + params.request_indices = nullptr; + params.qo_tile_indices = nullptr; + params.kv_tile_indices = nullptr; + params.merge_indptr = nullptr; + params.o_indptr = nullptr; + params.kv_chunk_size_ptr = nullptr; + params.block_valid_mask = nullptr; + params.total_num_rows = nullptr; + params.max_total_num_rows = 0; + params.padded_batch_size = 0; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + DTypeO* tmp_v = nullptr; + float* tmp_s = nullptr; + + params.request_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.request_indices_offset); + params.qo_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_tile_indices_offset); + params.kv_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_tile_indices_offset); + params.o_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.o_indptr_offset); + params.kv_chunk_size_ptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_chunk_size_ptr_offset); + if (plan_info.split_kv) { + params.merge_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.merge_indptr_offset); + tmp_v = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.v_offset); + tmp_s = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.s_offset); + if (plan_info.enable_cuda_graph) { + params.block_valid_mask = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.block_valid_mask_offset); + } + } + params.padded_batch_size = plan_info.padded_batch_size; + params.max_total_num_rows = plan_info.total_num_rows; + if (plan_info.enable_cuda_graph) { + params.total_num_rows = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.total_num_rows_offset); + } + + cudaError_t status = cudaSuccess; + + DISPATCH_CTA_TILE_Q(plan_info.cta_tile_q, CTA_TILE_Q, { + status = flashinfer::BatchPrefillWithPagedKVCacheDispatched< + CTA_TILE_Q, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + /*use_fp16_qk_reduction=*/USE_FP16_QK_REDUCTION, MASK_MODE, AttentionVariant, + PagedParams>(params, tmp_v, tmp_s, enable_pdl, stream); + }); + + TORCH_CHECK(status == cudaSuccess, "BatchPrefillWithPagedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); +} diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_config.inc b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..c43ae76fdd3ce360df6b8b0a7bb3ae43b142bc30 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_config.inc @@ -0,0 +1,153 @@ +#pragma once +#include +#include +#include +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_custom_mask, std::optional maybe_mask_indptr, std::optional maybe_alibi_slopes, std::optional maybe_prefix_len_ptr, std::optional maybe_token_pos_in_items_ptr, std::optional maybe_max_item_len_ptr, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta, int64_t token_pos_in_items_len +#define ADDITIONAL_PARAMS_SETTER params.maybe_custom_mask = maybe_custom_mask ? static_cast(maybe_custom_mask->data_ptr()): nullptr; \ +params.maybe_mask_indptr = maybe_mask_indptr ? static_cast(maybe_mask_indptr->data_ptr()): nullptr; \ +params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.maybe_prefix_len_ptr = maybe_prefix_len_ptr ? static_cast(maybe_prefix_len_ptr->data_ptr()): nullptr; \ +params.maybe_token_pos_in_items_ptr = maybe_token_pos_in_items_ptr ? static_cast(maybe_token_pos_in_items_ptr->data_ptr()): nullptr; \ +params.maybe_max_item_len_ptr = maybe_max_item_len_ptr ? static_cast(maybe_max_item_len_ptr->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; \ +params.token_pos_in_items_len = token_pos_in_items_len; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, RaggedParams, PagedParams, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { \ + constexpr auto use_custom_mask = MASK_MODE == MaskMode::kCustom; \ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ + }) + +using namespace flashinfer; + +using DTypeQ = nv_bfloat16; +using DTypeKV = __nv_fp8_e4m3; +using DTypeO = nv_bfloat16; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 64; +constexpr int HEAD_DIM_VO = 64; +constexpr bool USE_FP16_QK_REDUCTION = false; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = false; + + +struct RaggedParams { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + DTypeQ* q; + DTypeKV* k; + DTypeKV* v; + IdType* q_indptr; + IdType* kv_indptr; + DTypeO* o; + float* lse; + uint_fastdiv group_size; + + uint8_t* maybe_custom_mask; +int32_t* maybe_mask_indptr; +float* maybe_alibi_slopes; +uint32_t* maybe_prefix_len_ptr; +uint16_t* maybe_token_pos_in_items_ptr; +uint16_t* maybe_max_item_len_ptr; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; +int64_t token_pos_in_items_len; + + uint32_t num_qo_heads; + uint32_t num_kv_heads; + uint32_t q_stride_n; + uint32_t q_stride_h; + uint32_t k_stride_n; + uint32_t k_stride_h; + uint32_t v_stride_n; + uint32_t v_stride_h; + int32_t window_left; + + IdType* request_indices; + IdType* qo_tile_indices; + IdType* kv_tile_indices; + IdType* merge_indptr; + IdType* o_indptr; + IdType* kv_chunk_size_ptr; + bool* block_valid_mask; + uint32_t max_total_num_rows; + uint32_t* total_num_rows; + uint32_t padded_batch_size; + bool partition_kv; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { + return q_indptr[batch_idx + 1] - q_indptr[batch_idx]; + } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return kv_indptr[batch_idx + 1] - kv_indptr[batch_idx]; + } +}; + +struct PagedParams { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + DTypeQ* q; + paged_kv_t paged_kv; + IdType* q_indptr; + DTypeO* o; + float* lse; + uint_fastdiv group_size; + + uint8_t* maybe_custom_mask; +int32_t* maybe_mask_indptr; +float* maybe_alibi_slopes; +uint32_t* maybe_prefix_len_ptr; +uint16_t* maybe_token_pos_in_items_ptr; +uint16_t* maybe_max_item_len_ptr; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; +int64_t token_pos_in_items_len; + + uint32_t num_qo_heads; + IdType q_stride_n; + IdType q_stride_h; + int32_t window_left; + + IdType* request_indices; + IdType* qo_tile_indices; + IdType* kv_tile_indices; + IdType* merge_indptr; + IdType* o_indptr; + bool* block_valid_mask; + IdType* kv_chunk_size_ptr; + uint32_t max_total_num_rows; + uint32_t* total_num_rows; + uint32_t padded_batch_size; + bool partition_kv; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { + return q_indptr[batch_idx + 1] - q_indptr[batch_idx]; + } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return paged_kv.get_length(batch_idx); + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_jit_pybind.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..5421ab1cf910fdfbc867c182c8e22f23e44749ae --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_jit_pybind.cu @@ -0,0 +1,48 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "batch_prefill_config.inc" +#include "pytorch_extension_utils.h" + +at::Tensor BatchPrefillWithKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, + int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, + int64_t head_dim_vo, bool causal); + +void BatchPrefillWithRaggedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor k, at::Tensor v, + at::Tensor qo_indptr, at::Tensor kv_indptr, at::Tensor o, + std::optional maybe_lse, int64_t mask_mode_code, + int64_t layout, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +void BatchPrefillWithPagedKVCacheRun( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, + at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Batch-request prefill attention with KV-Cache plan + m.def("plan", BatchPrefillWithKVCachePlan); + // Batch-request prefill attention with KV-Cache operator + m.def("ragged_run", BatchPrefillWithRaggedKVCacheRun); + // Batch-request prefill attention with KV-Cache operator + m.def("paged_run", BatchPrefillWithPagedKVCacheRun); +} diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_0.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..dc470ac4930d3f74861b25a070ad59708802c8fd --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_0.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kNone == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_1.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..c94a2ee8dc6664dca2e0ffded85071bd6bd61864 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_1.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCausal == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_2.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..cfdc84f24991a8441236dd46f6a83b9dabe31531 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_2.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCustom == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_3.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..ea7d2bb0d4590de5f5d505a6f69f1654f4186ba6 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_3.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kMultiItemScoring == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_0.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..312ae24e057e022000309e4d9dab394e17b24613 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_0.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kNone == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_1.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..9df75df63e5639cedd81990d587d783f2fdc3efb --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_1.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCausal == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_2.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..476f7e93b1d1660b453d5c2d3dc7a4cacf23c24b --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_2.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCustom == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_3.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..ad1770813fab3d00b78cfec767f9df3475d31ea6 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_3.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kMultiItemScoring == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill.cu new file mode 100644 index 0000000000000000000000000000000000000000..a51fc7f56a0af7482bd5d2a9d3f672854df69dfd --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill.cu @@ -0,0 +1,330 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include +#include + +#include "batch_prefill_config.inc" +#include "pytorch_conversion_utils.h" +#include "pytorch_extension_utils.h" + +namespace flashinfer { + +template +cudaError_t BatchPrefillWithPagedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v, + float* tmp_s, bool enable_pdl, + cudaStream_t stream); + +template +cudaError_t BatchPrefillWithRaggedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v, + float* tmp_s, bool enable_pdl, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +at::Tensor BatchPrefillWithKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, + int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, + int64_t head_dim_vo, bool causal) { + size_t float_workspace_size_in_bytes = + float_workspace_buffer.size(0) * float_workspace_buffer.element_size(); + size_t int_workspace_size_in_bytes = + int_workspace_buffer.size(0) * int_workspace_buffer.element_size(); + + PrefillPlanInfo plan_info; + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + cudaError_t status = PrefillPlan( + float_workspace_buffer.data_ptr(), float_workspace_size_in_bytes, + int_workspace_buffer.data_ptr(), page_locked_int_workspace_buffer.data_ptr(), + int_workspace_size_in_bytes, plan_info, qo_indptr.data_ptr(), + kv_indptr.data_ptr(), total_num_rows, batch_size, num_qo_heads, num_kv_heads, + head_dim_qk, head_dim_vo, page_size, enable_cuda_graph, /*sizeof_dtype_o=*/2, stream); + + TORCH_CHECK(status == cudaSuccess, + "Failed to plan prefill with error: ", cudaGetErrorString(status)); + + return vec_to_tensor(plan_info.ToVector()); +} + +void BatchPrefillWithRaggedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor k, at::Tensor v, + at::Tensor qo_indptr, at::Tensor kv_indptr, at::Tensor o, + std::optional maybe_lse, int64_t mask_mode_code, + int64_t layout, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + PrefillPlanInfo plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + QKVLayout kv_layout = static_cast(layout); + + int64_t num_qo_heads = q.size(1); + int64_t head_dim_qk = q.size(2); + int64_t num_kv_heads = (kv_layout == QKVLayout::kNHD) ? k.size(1) : k.size(0); + uint32_t q_stride_n = q.stride(0), q_stride_h = q.stride(1), k_stride_n, k_stride_h, v_stride_n, + v_stride_h; + if (kv_layout == QKVLayout::kNHD) { + k_stride_n = k.stride(0); + k_stride_h = k.stride(1); + v_stride_n = v.stride(0); + v_stride_h = v.stride(1); + } else { + k_stride_h = k.stride(0); + k_stride_n = k.stride(1); + v_stride_h = v.stride(0); + v_stride_n = v.stride(1); + } + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == q.size(0), lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == q.size(1), lse.size(1), q.size(1)); + } + + void* float_buffer_ptr = float_workspace_buffer.data_ptr(); + void* int_buffer_ptr = int_workspace_buffer.data_ptr(); + + const MaskMode mask_mode = static_cast(mask_mode_code); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, + RaggedParams, PagedParams, [&] { + RaggedParams params; + + params.q = static_cast(q.data_ptr()); + params.k = static_cast(k.data_ptr()); + params.v = static_cast(v.data_ptr()); + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.q_indptr = static_cast(qo_indptr.data_ptr()); + params.kv_indptr = static_cast(kv_indptr.data_ptr()); + params.num_qo_heads = num_qo_heads; + params.num_kv_heads = num_kv_heads; + params.group_size = uint_fastdiv(num_qo_heads / num_kv_heads); + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.k_stride_n = k_stride_n; + params.k_stride_h = k_stride_h; + params.v_stride_n = v_stride_n; + params.v_stride_h = v_stride_h; + params.window_left = window_left; + + params.request_indices = nullptr; + params.qo_tile_indices = nullptr; + params.kv_tile_indices = nullptr; + params.merge_indptr = nullptr; + params.o_indptr = nullptr; + params.kv_chunk_size_ptr = nullptr; + params.block_valid_mask = nullptr; + params.total_num_rows = nullptr; + params.max_total_num_rows = 0; + params.padded_batch_size = 0; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + DTypeO* tmp_v = nullptr; + float* tmp_s = nullptr; + + params.request_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.request_indices_offset); + params.qo_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_tile_indices_offset); + params.kv_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_tile_indices_offset); + params.o_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.o_indptr_offset); + params.kv_chunk_size_ptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_chunk_size_ptr_offset); + if (plan_info.split_kv) { + params.merge_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.merge_indptr_offset); + tmp_v = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.v_offset); + tmp_s = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.s_offset); + if (plan_info.enable_cuda_graph) { + params.block_valid_mask = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.block_valid_mask_offset); + } + } + params.padded_batch_size = plan_info.padded_batch_size; + params.max_total_num_rows = plan_info.total_num_rows; + if (plan_info.enable_cuda_graph) { + params.total_num_rows = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.total_num_rows_offset); + } + + cudaError_t status = cudaSuccess; + + DISPATCH_CTA_TILE_Q(plan_info.cta_tile_q, CTA_TILE_Q, { + status = flashinfer::BatchPrefillWithRaggedKVCacheDispatched< + CTA_TILE_Q, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + /*use_fp16_qk_reduction=*/USE_FP16_QK_REDUCTION, MASK_MODE, AttentionVariant, + RaggedParams>(params, tmp_v, tmp_s, enable_pdl, stream); + }); + + TORCH_CHECK(status == cudaSuccess, "BatchPrefillWithRaggedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); +} + +void BatchPrefillWithPagedKVCacheRun( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, + at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + PrefillPlanInfo plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + QKVLayout kv_layout = static_cast(layout); + auto device = q.device(); + int64_t batch_size = paged_kv_indptr.size(0) - 1; + int64_t num_qo_heads = q.size(1); + int64_t num_kv_heads, page_size; + uint32_t head_dim_qk = q.size(2); + if (kv_layout == QKVLayout::kHND) { + num_kv_heads = paged_k_cache.size(1); + page_size = paged_k_cache.size(2); + } else { + page_size = paged_k_cache.size(1); + num_kv_heads = paged_k_cache.size(2); + } + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == q.size(0), lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == q.size(1), lse.size(1), q.size(1)); + } + + void* float_buffer_ptr = static_cast(float_workspace_buffer.data_ptr()); + void* int_buffer_ptr = static_cast(int_workspace_buffer.data_ptr()); + + const MaskMode mask_mode = static_cast(mask_mode_code); + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = paged_k_cache.scalar_type(); + + // get q_stride_n and q_stride_h + const auto q_stride_n = q.stride(0); + const auto q_stride_h = q.stride(1); + + // get kv_cache_strides + const int64_t* kv_cache_strides = nullptr; + auto k_strides = paged_k_cache.strides(); + auto v_strides = paged_v_cache.strides(); + TORCH_CHECK(k_strides == v_strides, "k/v strides must be identical"); + kv_cache_strides = k_strides.data(); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, + RaggedParams, PagedParams, [&] { + PagedParams params; + + params.q = static_cast(q.data_ptr()); + paged_kv_t paged_kv( + num_kv_heads, page_size, HEAD_DIM_VO, batch_size, kv_layout, + static_cast(paged_k_cache.data_ptr()), + static_cast(paged_v_cache.data_ptr()), kv_cache_strides, + static_cast(paged_kv_indices.data_ptr()), + static_cast(paged_kv_indptr.data_ptr()), + static_cast(paged_kv_last_page_len.data_ptr())); + params.paged_kv = paged_kv; + params.q_indptr = static_cast(qo_indptr.data_ptr()); + params.o = static_cast(o.data_ptr()); + + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.num_qo_heads = num_qo_heads; + params.group_size = uint_fastdiv(num_qo_heads / paged_kv.num_heads); + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.window_left = window_left; + + params.request_indices = nullptr; + params.qo_tile_indices = nullptr; + params.kv_tile_indices = nullptr; + params.merge_indptr = nullptr; + params.o_indptr = nullptr; + params.kv_chunk_size_ptr = nullptr; + params.block_valid_mask = nullptr; + params.total_num_rows = nullptr; + params.max_total_num_rows = 0; + params.padded_batch_size = 0; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + DTypeO* tmp_v = nullptr; + float* tmp_s = nullptr; + + params.request_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.request_indices_offset); + params.qo_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_tile_indices_offset); + params.kv_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_tile_indices_offset); + params.o_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.o_indptr_offset); + params.kv_chunk_size_ptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_chunk_size_ptr_offset); + if (plan_info.split_kv) { + params.merge_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.merge_indptr_offset); + tmp_v = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.v_offset); + tmp_s = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.s_offset); + if (plan_info.enable_cuda_graph) { + params.block_valid_mask = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.block_valid_mask_offset); + } + } + params.padded_batch_size = plan_info.padded_batch_size; + params.max_total_num_rows = plan_info.total_num_rows; + if (plan_info.enable_cuda_graph) { + params.total_num_rows = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.total_num_rows_offset); + } + + cudaError_t status = cudaSuccess; + + DISPATCH_CTA_TILE_Q(plan_info.cta_tile_q, CTA_TILE_Q, { + status = flashinfer::BatchPrefillWithPagedKVCacheDispatched< + CTA_TILE_Q, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + /*use_fp16_qk_reduction=*/USE_FP16_QK_REDUCTION, MASK_MODE, AttentionVariant, + PagedParams>(params, tmp_v, tmp_s, enable_pdl, stream); + }); + + TORCH_CHECK(status == cudaSuccess, "BatchPrefillWithPagedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); +} diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_config.inc b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..5284de65348f49eeb4a369650a649e227ede06bc --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_config.inc @@ -0,0 +1,153 @@ +#pragma once +#include +#include +#include +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_custom_mask, std::optional maybe_mask_indptr, std::optional maybe_alibi_slopes, std::optional maybe_prefix_len_ptr, std::optional maybe_token_pos_in_items_ptr, std::optional maybe_max_item_len_ptr, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta, int64_t token_pos_in_items_len +#define ADDITIONAL_PARAMS_SETTER params.maybe_custom_mask = maybe_custom_mask ? static_cast(maybe_custom_mask->data_ptr()): nullptr; \ +params.maybe_mask_indptr = maybe_mask_indptr ? static_cast(maybe_mask_indptr->data_ptr()): nullptr; \ +params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.maybe_prefix_len_ptr = maybe_prefix_len_ptr ? static_cast(maybe_prefix_len_ptr->data_ptr()): nullptr; \ +params.maybe_token_pos_in_items_ptr = maybe_token_pos_in_items_ptr ? static_cast(maybe_token_pos_in_items_ptr->data_ptr()): nullptr; \ +params.maybe_max_item_len_ptr = maybe_max_item_len_ptr ? static_cast(maybe_max_item_len_ptr->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; \ +params.token_pos_in_items_len = token_pos_in_items_len; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, RaggedParams, PagedParams, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { \ + constexpr auto use_custom_mask = MASK_MODE == MaskMode::kCustom; \ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ + }) + +using namespace flashinfer; + +using DTypeQ = nv_bfloat16; +using DTypeKV = __nv_fp8_e4m3; +using DTypeO = nv_bfloat16; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 64; +constexpr int HEAD_DIM_VO = 64; +constexpr bool USE_FP16_QK_REDUCTION = false; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = true; + + +struct RaggedParams { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + DTypeQ* q; + DTypeKV* k; + DTypeKV* v; + IdType* q_indptr; + IdType* kv_indptr; + DTypeO* o; + float* lse; + uint_fastdiv group_size; + + uint8_t* maybe_custom_mask; +int32_t* maybe_mask_indptr; +float* maybe_alibi_slopes; +uint32_t* maybe_prefix_len_ptr; +uint16_t* maybe_token_pos_in_items_ptr; +uint16_t* maybe_max_item_len_ptr; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; +int64_t token_pos_in_items_len; + + uint32_t num_qo_heads; + uint32_t num_kv_heads; + uint32_t q_stride_n; + uint32_t q_stride_h; + uint32_t k_stride_n; + uint32_t k_stride_h; + uint32_t v_stride_n; + uint32_t v_stride_h; + int32_t window_left; + + IdType* request_indices; + IdType* qo_tile_indices; + IdType* kv_tile_indices; + IdType* merge_indptr; + IdType* o_indptr; + IdType* kv_chunk_size_ptr; + bool* block_valid_mask; + uint32_t max_total_num_rows; + uint32_t* total_num_rows; + uint32_t padded_batch_size; + bool partition_kv; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { + return q_indptr[batch_idx + 1] - q_indptr[batch_idx]; + } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return kv_indptr[batch_idx + 1] - kv_indptr[batch_idx]; + } +}; + +struct PagedParams { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + DTypeQ* q; + paged_kv_t paged_kv; + IdType* q_indptr; + DTypeO* o; + float* lse; + uint_fastdiv group_size; + + uint8_t* maybe_custom_mask; +int32_t* maybe_mask_indptr; +float* maybe_alibi_slopes; +uint32_t* maybe_prefix_len_ptr; +uint16_t* maybe_token_pos_in_items_ptr; +uint16_t* maybe_max_item_len_ptr; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; +int64_t token_pos_in_items_len; + + uint32_t num_qo_heads; + IdType q_stride_n; + IdType q_stride_h; + int32_t window_left; + + IdType* request_indices; + IdType* qo_tile_indices; + IdType* kv_tile_indices; + IdType* merge_indptr; + IdType* o_indptr; + bool* block_valid_mask; + IdType* kv_chunk_size_ptr; + uint32_t max_total_num_rows; + uint32_t* total_num_rows; + uint32_t padded_batch_size; + bool partition_kv; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { + return q_indptr[batch_idx + 1] - q_indptr[batch_idx]; + } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return paged_kv.get_length(batch_idx); + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_jit_pybind.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..5421ab1cf910fdfbc867c182c8e22f23e44749ae --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_jit_pybind.cu @@ -0,0 +1,48 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "batch_prefill_config.inc" +#include "pytorch_extension_utils.h" + +at::Tensor BatchPrefillWithKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, + int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, + int64_t head_dim_vo, bool causal); + +void BatchPrefillWithRaggedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor k, at::Tensor v, + at::Tensor qo_indptr, at::Tensor kv_indptr, at::Tensor o, + std::optional maybe_lse, int64_t mask_mode_code, + int64_t layout, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +void BatchPrefillWithPagedKVCacheRun( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, + at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Batch-request prefill attention with KV-Cache plan + m.def("plan", BatchPrefillWithKVCachePlan); + // Batch-request prefill attention with KV-Cache operator + m.def("ragged_run", BatchPrefillWithRaggedKVCacheRun); + // Batch-request prefill attention with KV-Cache operator + m.def("paged_run", BatchPrefillWithPagedKVCacheRun); +} diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_0.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..e18cf3fb9d174fdedc0fd1a8b54cd9bdba181250 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_0.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kNone == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_1.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..2ca8d62b9a6503e0367aa85a39fc5cc164a2e1be --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_1.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCausal == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_2.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..4f16cdb3774a6c0b5727ae57ad2672298a4e1446 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_2.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCustom == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_3.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..5143b75b7cc8d697393aa524693ff750f92bcb82 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_3.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kMultiItemScoring == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, PagedParams>(PagedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_0.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..21713a0ad335444b8200c6481c5b8d4619082752 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_0.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kNone == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_1.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..503fe49906ece459dfe133fc9ee55a1a3ace4625 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_1.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCausal == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_2.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..edcf2c0a482136da1804d7e6d038946b01cb2d49 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_2.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCustom == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_3.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..a2419d83e0324f8b398b949e02a95744b1a64e3b --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_3.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kMultiItemScoring == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, RaggedParams>(RaggedParams params, nv_bfloat16* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_fp8_sm90.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_fp8_sm90.cu new file mode 100644 index 0000000000000000000000000000000000000000..9f356754a8108a2e095d971c14a93ef79f7cfb4a --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_fp8_sm90.cu @@ -0,0 +1,185 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include +#include + +#include "batch_prefill_sm90_config.inc" +#include "pytorch_conversion_utils.h" +#include "pytorch_extension_utils.h" + +namespace flashinfer { + +template +cudaError_t BatchFP8PrefillWithPagedKVCacheDispatched(Params& params, bool enable_pdl, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +at::Tensor BatchPrefillWithKVCacheSM90Plan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, + int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, + int64_t head_dim_vo, bool causal) { + size_t float_workspace_size_in_bytes = + float_workspace_buffer.size(0) * float_workspace_buffer.element_size(); + size_t int_workspace_size_in_bytes = + int_workspace_buffer.size(0) * int_workspace_buffer.element_size(); + + flashinfer::PrefillPlanSM90Info plan_info; + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + cudaError_t status = + PrefillSM90Plan(float_workspace_buffer.data_ptr(), float_workspace_size_in_bytes, + int_workspace_buffer.data_ptr(), page_locked_int_workspace_buffer.data_ptr(), + int_workspace_size_in_bytes, plan_info, qo_indptr.data_ptr(), + kv_indptr.data_ptr(), kv_len_arr.data_ptr(), total_num_rows, + batch_size, num_qo_heads, num_kv_heads, head_dim_qk, head_dim_vo, page_size, + causal, enable_cuda_graph, /*sizeof_dtype_o=*/2, stream); + + TORCH_CHECK(status == cudaSuccess, + "PrefillSM90Plan failed with error: ", cudaGetErrorString(status)); + + return vec_to_tensor(plan_info.ToVector()); +} + +void BatchPrefillWithRaggedKVCacheSM90Run(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor k, at::Tensor v, + at::Tensor qo_indptr, at::Tensor kv_indptr, at::Tensor o, + std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left, + bool enable_pdl // placeholder + ADDITIONAL_FUNC_PARAMS) { + return; // TODO: Implement this function +} + +void BatchPrefillWithPagedKVCacheSM90Run( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, + at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + PrefillPlanSM90Info plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == q.size(0), lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == q.size(1), lse.size(1), q.size(1)); + } + QKVLayout kv_layout = static_cast(layout); + int64_t num_kv_heads, page_size; + int64_t head_dim_qk = q.size(2); + int64_t head_dim_vo = paged_v_cache.size(3); + if (kv_layout == QKVLayout::kHND) { + num_kv_heads = paged_k_cache.size(1); + page_size = paged_k_cache.size(2); + } else { + page_size = paged_k_cache.size(1); + num_kv_heads = paged_k_cache.size(2); + } + + void* float_buffer_ptr = float_workspace_buffer.data_ptr(); + void* int_buffer_ptr = int_workspace_buffer.data_ptr(); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = paged_k_cache.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + const MaskMode mask_mode = static_cast(mask_mode_code); + bool use_swa = window_left != -1; + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, USE_SLIDING_WINDOW, + USE_LOGITS_SOFT_CAP, AttentionVariant, RaggedParams, PagedParams, [&] { + PagedParams params; + + params.q_ptr = static_cast(q.data_ptr()); + params.k_ptr = static_cast(paged_k_cache.data_ptr()); + params.v_ptr = static_cast(paged_v_cache.data_ptr()); + params.o_ptr = static_cast(o.data_ptr()); + params.lse_ptr = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.q_stride_n = q.stride(0); + params.q_stride_h = q.stride(1); + params.o_stride_n = o.stride(0); + params.o_stride_h = o.stride(1); + if (kv_layout == QKVLayout::kNHD) { + // (num_pages, page_size, num_heads, head_dim) + params.k_stride_n = paged_k_cache.stride(1); + params.k_stride_h = paged_k_cache.stride(2); + params.v_stride_n = paged_v_cache.stride(1); + params.v_stride_h = paged_v_cache.stride(2); + } else { + // (num_pages, num_heads, page_size, head_dim) + params.k_stride_h = paged_k_cache.stride(1); + params.k_stride_n = paged_k_cache.stride(2); + params.v_stride_h = paged_v_cache.stride(1); + params.v_stride_n = paged_v_cache.stride(2); + } + params.nnz_qo = q.size(0); + params.num_qo_heads = q.size(1); + params.num_kv_heads = num_kv_heads; + params.group_size = params.num_qo_heads / num_kv_heads; + params.page_size = page_size; + params.window_left = window_left; + params.causal = mask_mode_code == 1; + params.qo_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_tile_indices_offset); + params.qo_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_indptr_offset); + params.kv_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_indptr_offset); + params.qo_lens = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_len_offset); + params.kv_lens = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_len_offset); + params.batch_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.batch_indices_offset); + params.head_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.head_indices_offset); + params.work_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.work_indptr_offset); + params.kv_indices = static_cast(paged_kv_indices.data_ptr()); + + ADDITIONAL_PARAMS_SETTER + + // Not support various head_dim for now + static_assert(HEAD_DIM_QK == HEAD_DIM_VO, "head_dim_qk and head_dim_vo should be the same"); + // Currently only support same quantization precision + static_assert(std::is_same_v); + + bool same_schedule_for_all_heads = plan_info.same_schedule_for_all_heads; + DISPATCH_BOOL(same_schedule_for_all_heads, SAME_SCHEDULER_FOR_ALL_HEADS, [&] { + cudaError_t status = + BatchFP8PrefillWithPagedKVCacheDispatched(params, enable_pdl, + stream); + TORCH_CHECK(status == cudaSuccess, + "BatchPrefillWithPagedKVCacheSM90Run failed with error: ", + cudaGetErrorString(status)); + return true; + }); + }); +} diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_0.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..db8cade58529f1f4915e3a3b2d6900719bcfdb44 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_0.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchFP8PrefillWithPagedKVCacheDispatched + <128, + MaskMode::kNone, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultFP8Attention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchFP8PrefillWithPagedKVCacheDispatched + <128, + MaskMode::kNone, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultFP8Attention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_1.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..c0982d059808c912ee27445f54fb4630011f9d46 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_1.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchFP8PrefillWithPagedKVCacheDispatched + <128, + MaskMode::kCausal, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultFP8Attention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchFP8PrefillWithPagedKVCacheDispatched + <128, + MaskMode::kCausal, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultFP8Attention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_2.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..1116f67965b40e78d68b33f975b141cbd98159fa --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_2.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchFP8PrefillWithPagedKVCacheDispatched + <128, + MaskMode::kCustom, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultFP8Attention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchFP8PrefillWithPagedKVCacheDispatched + <128, + MaskMode::kCustom, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultFP8Attention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_3.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..629cdbf5a1bdea4b6c380a6b4f150e58913b3548 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_3.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchFP8PrefillWithPagedKVCacheDispatched + <128, + MaskMode::kMultiItemScoring, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultFP8Attention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchFP8PrefillWithPagedKVCacheDispatched + <128, + MaskMode::kMultiItemScoring, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultFP8Attention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_0.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..730347d4abcaff413b5286e36cc0dfe09228751d --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_0.cu @@ -0,0 +1 @@ +// TODO: Not implemented yet \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_1.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..730347d4abcaff413b5286e36cc0dfe09228751d --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_1.cu @@ -0,0 +1 @@ +// TODO: Not implemented yet \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_2.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..730347d4abcaff413b5286e36cc0dfe09228751d --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_2.cu @@ -0,0 +1 @@ +// TODO: Not implemented yet \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_3.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..730347d4abcaff413b5286e36cc0dfe09228751d --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_3.cu @@ -0,0 +1 @@ +// TODO: Not implemented yet \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_sm90_config.inc b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_sm90_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..96d6b95598129179077ae12b1bfe22c7fa37a116 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_sm90_config.inc @@ -0,0 +1,129 @@ +#pragma once +#include +#include +#include +#include +#include +#include "pytorch_extension_utils.h" + +#define ADDITIONAL_FUNC_PARAMS , at::Tensor scale_q, at::Tensor scale_k, at::Tensor scale_v, double sm_scale +#define ADDITIONAL_PARAMS_SETTER params.additional_params.scale_q = static_cast(scale_q.data_ptr()); \ +params.additional_params.scale_k = static_cast(scale_k.data_ptr()); \ +params.additional_params.scale_v = static_cast(scale_v.data_ptr()); \ +params.additional_params.sm_scale = sm_scale; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, RaggedParams, PagedParams, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { using AttentionVariant = DefaultFP8Attention; __VA_ARGS__();}) + +using namespace flashinfer; + +using DTypeQ = cutlass_dtype_t<__nv_fp8_e4m3>; +using DTypeKV = cutlass_dtype_t<__nv_fp8_e4m3>; +using DTypeO = cutlass_dtype_t; +using IdType = cutlass_dtype_t; + +constexpr int HEAD_DIM_QK = 128; +constexpr int HEAD_DIM_VO = 128; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto USE_SLIDING_WINDOW = false; + +struct RaggedParams { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + // The QKV matrices. + DTypeQ* q_ptr; + DTypeKV* k_ptr; + DTypeKV* v_ptr; + DTypeO* o_ptr; + float* lse_ptr; + + IdType* qo_tile_indices; + IdType* qo_indptr; + IdType* kv_indptr; + IdType* qo_lens; + IdType* kv_lens; + IdType* head_indices; + IdType* work_indptr; + IdType* batch_indices; + + struct AdditionalParams { + float* scale_q; +float* scale_k; +float* scale_v; +double sm_scale; + + } additional_params; + + int64_t q_stride_n; + int64_t k_stride_n; + int64_t v_stride_n; + int64_t o_stride_n; + int64_t q_stride_h; + int64_t k_stride_h; + int64_t v_stride_h; + int64_t o_stride_h; + int64_t nnz_qo; + int64_t nnz_kv; + + int head_dim; + int num_qo_heads; + int num_kv_heads; + int group_size; + int window_left; + + bool causal; +}; + +struct PagedParams { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + // The QKV matrices. + DTypeQ* q_ptr; + DTypeKV* k_ptr; + DTypeKV* v_ptr; + DTypeO* o_ptr; + float* lse_ptr; + + IdType* qo_tile_indices; + IdType* qo_indptr; + IdType* kv_indptr; + IdType* kv_indices; + IdType* qo_lens; + IdType* kv_lens; + IdType* head_indices; + IdType* work_indptr; + IdType* batch_indices; + + struct AdditionalParams { + float* scale_q; +float* scale_k; +float* scale_v; +double sm_scale; + + } additional_params; + + int64_t q_stride_n; + int64_t k_stride_n; + int64_t v_stride_n; + int64_t o_stride_n; + int64_t q_stride_h; + int64_t k_stride_h; + int64_t v_stride_h; + int64_t o_stride_h; + int64_t nnz_qo; + + int head_dim; + int num_qo_heads; + int num_kv_heads; + int group_size; + int page_size; + int window_left; + + bool causal; +}; + +#include \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_sm90_jit_pybind.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_sm90_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..b1edd582549e124c776f2cf9b4503f4bf60d16ae --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_sm90_jit_pybind.cu @@ -0,0 +1,46 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "batch_prefill_sm90_config.inc" +#include "pytorch_extension_utils.h" + +at::Tensor BatchPrefillWithKVCacheSM90Plan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, + int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, + int64_t head_dim_vo, bool causal); + +void BatchPrefillWithRaggedKVCacheSM90Run( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +void BatchPrefillWithPagedKVCacheSM90Run( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, + at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Batch-request prefill attention with KV-Cache plan + m.def("plan", BatchPrefillWithKVCacheSM90Plan); + // Batch-request prefill attention with KV-Cache operator + m.def("ragged_run", BatchPrefillWithRaggedKVCacheSM90Run); + // Batch-request prefill attention with KV-Cache operator + m.def("paged_run", BatchPrefillWithPagedKVCacheSM90Run); +} diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_fp8_sm90.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_fp8_sm90.cu new file mode 100644 index 0000000000000000000000000000000000000000..9f356754a8108a2e095d971c14a93ef79f7cfb4a --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_fp8_sm90.cu @@ -0,0 +1,185 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include +#include + +#include "batch_prefill_sm90_config.inc" +#include "pytorch_conversion_utils.h" +#include "pytorch_extension_utils.h" + +namespace flashinfer { + +template +cudaError_t BatchFP8PrefillWithPagedKVCacheDispatched(Params& params, bool enable_pdl, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +at::Tensor BatchPrefillWithKVCacheSM90Plan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, + int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, + int64_t head_dim_vo, bool causal) { + size_t float_workspace_size_in_bytes = + float_workspace_buffer.size(0) * float_workspace_buffer.element_size(); + size_t int_workspace_size_in_bytes = + int_workspace_buffer.size(0) * int_workspace_buffer.element_size(); + + flashinfer::PrefillPlanSM90Info plan_info; + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + cudaError_t status = + PrefillSM90Plan(float_workspace_buffer.data_ptr(), float_workspace_size_in_bytes, + int_workspace_buffer.data_ptr(), page_locked_int_workspace_buffer.data_ptr(), + int_workspace_size_in_bytes, plan_info, qo_indptr.data_ptr(), + kv_indptr.data_ptr(), kv_len_arr.data_ptr(), total_num_rows, + batch_size, num_qo_heads, num_kv_heads, head_dim_qk, head_dim_vo, page_size, + causal, enable_cuda_graph, /*sizeof_dtype_o=*/2, stream); + + TORCH_CHECK(status == cudaSuccess, + "PrefillSM90Plan failed with error: ", cudaGetErrorString(status)); + + return vec_to_tensor(plan_info.ToVector()); +} + +void BatchPrefillWithRaggedKVCacheSM90Run(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor k, at::Tensor v, + at::Tensor qo_indptr, at::Tensor kv_indptr, at::Tensor o, + std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left, + bool enable_pdl // placeholder + ADDITIONAL_FUNC_PARAMS) { + return; // TODO: Implement this function +} + +void BatchPrefillWithPagedKVCacheSM90Run( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, + at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + PrefillPlanSM90Info plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == q.size(0), lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == q.size(1), lse.size(1), q.size(1)); + } + QKVLayout kv_layout = static_cast(layout); + int64_t num_kv_heads, page_size; + int64_t head_dim_qk = q.size(2); + int64_t head_dim_vo = paged_v_cache.size(3); + if (kv_layout == QKVLayout::kHND) { + num_kv_heads = paged_k_cache.size(1); + page_size = paged_k_cache.size(2); + } else { + page_size = paged_k_cache.size(1); + num_kv_heads = paged_k_cache.size(2); + } + + void* float_buffer_ptr = float_workspace_buffer.data_ptr(); + void* int_buffer_ptr = int_workspace_buffer.data_ptr(); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = paged_k_cache.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + const MaskMode mask_mode = static_cast(mask_mode_code); + bool use_swa = window_left != -1; + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, USE_SLIDING_WINDOW, + USE_LOGITS_SOFT_CAP, AttentionVariant, RaggedParams, PagedParams, [&] { + PagedParams params; + + params.q_ptr = static_cast(q.data_ptr()); + params.k_ptr = static_cast(paged_k_cache.data_ptr()); + params.v_ptr = static_cast(paged_v_cache.data_ptr()); + params.o_ptr = static_cast(o.data_ptr()); + params.lse_ptr = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.q_stride_n = q.stride(0); + params.q_stride_h = q.stride(1); + params.o_stride_n = o.stride(0); + params.o_stride_h = o.stride(1); + if (kv_layout == QKVLayout::kNHD) { + // (num_pages, page_size, num_heads, head_dim) + params.k_stride_n = paged_k_cache.stride(1); + params.k_stride_h = paged_k_cache.stride(2); + params.v_stride_n = paged_v_cache.stride(1); + params.v_stride_h = paged_v_cache.stride(2); + } else { + // (num_pages, num_heads, page_size, head_dim) + params.k_stride_h = paged_k_cache.stride(1); + params.k_stride_n = paged_k_cache.stride(2); + params.v_stride_h = paged_v_cache.stride(1); + params.v_stride_n = paged_v_cache.stride(2); + } + params.nnz_qo = q.size(0); + params.num_qo_heads = q.size(1); + params.num_kv_heads = num_kv_heads; + params.group_size = params.num_qo_heads / num_kv_heads; + params.page_size = page_size; + params.window_left = window_left; + params.causal = mask_mode_code == 1; + params.qo_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_tile_indices_offset); + params.qo_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_indptr_offset); + params.kv_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_indptr_offset); + params.qo_lens = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_len_offset); + params.kv_lens = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_len_offset); + params.batch_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.batch_indices_offset); + params.head_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.head_indices_offset); + params.work_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.work_indptr_offset); + params.kv_indices = static_cast(paged_kv_indices.data_ptr()); + + ADDITIONAL_PARAMS_SETTER + + // Not support various head_dim for now + static_assert(HEAD_DIM_QK == HEAD_DIM_VO, "head_dim_qk and head_dim_vo should be the same"); + // Currently only support same quantization precision + static_assert(std::is_same_v); + + bool same_schedule_for_all_heads = plan_info.same_schedule_for_all_heads; + DISPATCH_BOOL(same_schedule_for_all_heads, SAME_SCHEDULER_FOR_ALL_HEADS, [&] { + cudaError_t status = + BatchFP8PrefillWithPagedKVCacheDispatched(params, enable_pdl, + stream); + TORCH_CHECK(status == cudaSuccess, + "BatchPrefillWithPagedKVCacheSM90Run failed with error: ", + cudaGetErrorString(status)); + return true; + }); + }); +} diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_0.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..12fe96fede547f29a056b9925ee32001e82f49fd --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_0.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchFP8PrefillWithPagedKVCacheDispatched + <256, + MaskMode::kNone, + /*USE_SLIDING_WINDOW=*/true, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultFP8Attention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchFP8PrefillWithPagedKVCacheDispatched + <256, + MaskMode::kNone, + /*USE_SLIDING_WINDOW=*/true, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultFP8Attention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_1.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..53de9f1d9b0288a3a35245f794c09ab084a06271 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_1.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchFP8PrefillWithPagedKVCacheDispatched + <256, + MaskMode::kCausal, + /*USE_SLIDING_WINDOW=*/true, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultFP8Attention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchFP8PrefillWithPagedKVCacheDispatched + <256, + MaskMode::kCausal, + /*USE_SLIDING_WINDOW=*/true, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultFP8Attention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_2.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..d3854bf6989d9916f5045f1cc77b0504547219fa --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_2.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchFP8PrefillWithPagedKVCacheDispatched + <256, + MaskMode::kCustom, + /*USE_SLIDING_WINDOW=*/true, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultFP8Attention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchFP8PrefillWithPagedKVCacheDispatched + <256, + MaskMode::kCustom, + /*USE_SLIDING_WINDOW=*/true, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultFP8Attention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_3.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..fac46d2f0439c7c36c9b95169b3b864a86c7bb13 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_3.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchFP8PrefillWithPagedKVCacheDispatched + <256, + MaskMode::kMultiItemScoring, + /*USE_SLIDING_WINDOW=*/true, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultFP8Attention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchFP8PrefillWithPagedKVCacheDispatched + <256, + MaskMode::kMultiItemScoring, + /*USE_SLIDING_WINDOW=*/true, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultFP8Attention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_0.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..730347d4abcaff413b5286e36cc0dfe09228751d --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_0.cu @@ -0,0 +1 @@ +// TODO: Not implemented yet \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_1.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..730347d4abcaff413b5286e36cc0dfe09228751d --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_1.cu @@ -0,0 +1 @@ +// TODO: Not implemented yet \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_2.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..730347d4abcaff413b5286e36cc0dfe09228751d --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_2.cu @@ -0,0 +1 @@ +// TODO: Not implemented yet \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_3.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..730347d4abcaff413b5286e36cc0dfe09228751d --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_3.cu @@ -0,0 +1 @@ +// TODO: Not implemented yet \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_sm90_config.inc b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_sm90_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..674e9fb3a905205aa5a45c8b4fab4d30921d4eb4 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_sm90_config.inc @@ -0,0 +1,129 @@ +#pragma once +#include +#include +#include +#include +#include +#include "pytorch_extension_utils.h" + +#define ADDITIONAL_FUNC_PARAMS , at::Tensor scale_q, at::Tensor scale_k, at::Tensor scale_v, double sm_scale +#define ADDITIONAL_PARAMS_SETTER params.additional_params.scale_q = static_cast(scale_q.data_ptr()); \ +params.additional_params.scale_k = static_cast(scale_k.data_ptr()); \ +params.additional_params.scale_v = static_cast(scale_v.data_ptr()); \ +params.additional_params.sm_scale = sm_scale; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, RaggedParams, PagedParams, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { using AttentionVariant = DefaultFP8Attention; __VA_ARGS__();}) + +using namespace flashinfer; + +using DTypeQ = cutlass_dtype_t<__nv_fp8_e4m3>; +using DTypeKV = cutlass_dtype_t<__nv_fp8_e4m3>; +using DTypeO = cutlass_dtype_t; +using IdType = cutlass_dtype_t; + +constexpr int HEAD_DIM_QK = 256; +constexpr int HEAD_DIM_VO = 256; +constexpr auto USE_LOGITS_SOFT_CAP = true; +constexpr auto USE_SLIDING_WINDOW = true; + +struct RaggedParams { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + // The QKV matrices. + DTypeQ* q_ptr; + DTypeKV* k_ptr; + DTypeKV* v_ptr; + DTypeO* o_ptr; + float* lse_ptr; + + IdType* qo_tile_indices; + IdType* qo_indptr; + IdType* kv_indptr; + IdType* qo_lens; + IdType* kv_lens; + IdType* head_indices; + IdType* work_indptr; + IdType* batch_indices; + + struct AdditionalParams { + float* scale_q; +float* scale_k; +float* scale_v; +double sm_scale; + + } additional_params; + + int64_t q_stride_n; + int64_t k_stride_n; + int64_t v_stride_n; + int64_t o_stride_n; + int64_t q_stride_h; + int64_t k_stride_h; + int64_t v_stride_h; + int64_t o_stride_h; + int64_t nnz_qo; + int64_t nnz_kv; + + int head_dim; + int num_qo_heads; + int num_kv_heads; + int group_size; + int window_left; + + bool causal; +}; + +struct PagedParams { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + // The QKV matrices. + DTypeQ* q_ptr; + DTypeKV* k_ptr; + DTypeKV* v_ptr; + DTypeO* o_ptr; + float* lse_ptr; + + IdType* qo_tile_indices; + IdType* qo_indptr; + IdType* kv_indptr; + IdType* kv_indices; + IdType* qo_lens; + IdType* kv_lens; + IdType* head_indices; + IdType* work_indptr; + IdType* batch_indices; + + struct AdditionalParams { + float* scale_q; +float* scale_k; +float* scale_v; +double sm_scale; + + } additional_params; + + int64_t q_stride_n; + int64_t k_stride_n; + int64_t v_stride_n; + int64_t o_stride_n; + int64_t q_stride_h; + int64_t k_stride_h; + int64_t v_stride_h; + int64_t o_stride_h; + int64_t nnz_qo; + + int head_dim; + int num_qo_heads; + int num_kv_heads; + int group_size; + int page_size; + int window_left; + + bool causal; +}; + +#include \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_sm90_jit_pybind.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_sm90_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..b1edd582549e124c776f2cf9b4503f4bf60d16ae --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_sm90_jit_pybind.cu @@ -0,0 +1,46 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "batch_prefill_sm90_config.inc" +#include "pytorch_extension_utils.h" + +at::Tensor BatchPrefillWithKVCacheSM90Plan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, + int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, + int64_t head_dim_vo, bool causal); + +void BatchPrefillWithRaggedKVCacheSM90Run( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +void BatchPrefillWithPagedKVCacheSM90Run( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, + at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Batch-request prefill attention with KV-Cache plan + m.def("plan", BatchPrefillWithKVCacheSM90Plan); + // Batch-request prefill attention with KV-Cache operator + m.def("ragged_run", BatchPrefillWithRaggedKVCacheSM90Run); + // Batch-request prefill attention with KV-Cache operator + m.def("paged_run", BatchPrefillWithPagedKVCacheSM90Run); +} diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_fp8_sm90.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_fp8_sm90.cu new file mode 100644 index 0000000000000000000000000000000000000000..9f356754a8108a2e095d971c14a93ef79f7cfb4a --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_fp8_sm90.cu @@ -0,0 +1,185 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include +#include + +#include "batch_prefill_sm90_config.inc" +#include "pytorch_conversion_utils.h" +#include "pytorch_extension_utils.h" + +namespace flashinfer { + +template +cudaError_t BatchFP8PrefillWithPagedKVCacheDispatched(Params& params, bool enable_pdl, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +at::Tensor BatchPrefillWithKVCacheSM90Plan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, + int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, + int64_t head_dim_vo, bool causal) { + size_t float_workspace_size_in_bytes = + float_workspace_buffer.size(0) * float_workspace_buffer.element_size(); + size_t int_workspace_size_in_bytes = + int_workspace_buffer.size(0) * int_workspace_buffer.element_size(); + + flashinfer::PrefillPlanSM90Info plan_info; + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + cudaError_t status = + PrefillSM90Plan(float_workspace_buffer.data_ptr(), float_workspace_size_in_bytes, + int_workspace_buffer.data_ptr(), page_locked_int_workspace_buffer.data_ptr(), + int_workspace_size_in_bytes, plan_info, qo_indptr.data_ptr(), + kv_indptr.data_ptr(), kv_len_arr.data_ptr(), total_num_rows, + batch_size, num_qo_heads, num_kv_heads, head_dim_qk, head_dim_vo, page_size, + causal, enable_cuda_graph, /*sizeof_dtype_o=*/2, stream); + + TORCH_CHECK(status == cudaSuccess, + "PrefillSM90Plan failed with error: ", cudaGetErrorString(status)); + + return vec_to_tensor(plan_info.ToVector()); +} + +void BatchPrefillWithRaggedKVCacheSM90Run(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor k, at::Tensor v, + at::Tensor qo_indptr, at::Tensor kv_indptr, at::Tensor o, + std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left, + bool enable_pdl // placeholder + ADDITIONAL_FUNC_PARAMS) { + return; // TODO: Implement this function +} + +void BatchPrefillWithPagedKVCacheSM90Run( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, + at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + PrefillPlanSM90Info plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == q.size(0), lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == q.size(1), lse.size(1), q.size(1)); + } + QKVLayout kv_layout = static_cast(layout); + int64_t num_kv_heads, page_size; + int64_t head_dim_qk = q.size(2); + int64_t head_dim_vo = paged_v_cache.size(3); + if (kv_layout == QKVLayout::kHND) { + num_kv_heads = paged_k_cache.size(1); + page_size = paged_k_cache.size(2); + } else { + page_size = paged_k_cache.size(1); + num_kv_heads = paged_k_cache.size(2); + } + + void* float_buffer_ptr = float_workspace_buffer.data_ptr(); + void* int_buffer_ptr = int_workspace_buffer.data_ptr(); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = paged_k_cache.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + const MaskMode mask_mode = static_cast(mask_mode_code); + bool use_swa = window_left != -1; + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, USE_SLIDING_WINDOW, + USE_LOGITS_SOFT_CAP, AttentionVariant, RaggedParams, PagedParams, [&] { + PagedParams params; + + params.q_ptr = static_cast(q.data_ptr()); + params.k_ptr = static_cast(paged_k_cache.data_ptr()); + params.v_ptr = static_cast(paged_v_cache.data_ptr()); + params.o_ptr = static_cast(o.data_ptr()); + params.lse_ptr = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.q_stride_n = q.stride(0); + params.q_stride_h = q.stride(1); + params.o_stride_n = o.stride(0); + params.o_stride_h = o.stride(1); + if (kv_layout == QKVLayout::kNHD) { + // (num_pages, page_size, num_heads, head_dim) + params.k_stride_n = paged_k_cache.stride(1); + params.k_stride_h = paged_k_cache.stride(2); + params.v_stride_n = paged_v_cache.stride(1); + params.v_stride_h = paged_v_cache.stride(2); + } else { + // (num_pages, num_heads, page_size, head_dim) + params.k_stride_h = paged_k_cache.stride(1); + params.k_stride_n = paged_k_cache.stride(2); + params.v_stride_h = paged_v_cache.stride(1); + params.v_stride_n = paged_v_cache.stride(2); + } + params.nnz_qo = q.size(0); + params.num_qo_heads = q.size(1); + params.num_kv_heads = num_kv_heads; + params.group_size = params.num_qo_heads / num_kv_heads; + params.page_size = page_size; + params.window_left = window_left; + params.causal = mask_mode_code == 1; + params.qo_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_tile_indices_offset); + params.qo_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_indptr_offset); + params.kv_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_indptr_offset); + params.qo_lens = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_len_offset); + params.kv_lens = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_len_offset); + params.batch_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.batch_indices_offset); + params.head_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.head_indices_offset); + params.work_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.work_indptr_offset); + params.kv_indices = static_cast(paged_kv_indices.data_ptr()); + + ADDITIONAL_PARAMS_SETTER + + // Not support various head_dim for now + static_assert(HEAD_DIM_QK == HEAD_DIM_VO, "head_dim_qk and head_dim_vo should be the same"); + // Currently only support same quantization precision + static_assert(std::is_same_v); + + bool same_schedule_for_all_heads = plan_info.same_schedule_for_all_heads; + DISPATCH_BOOL(same_schedule_for_all_heads, SAME_SCHEDULER_FOR_ALL_HEADS, [&] { + cudaError_t status = + BatchFP8PrefillWithPagedKVCacheDispatched(params, enable_pdl, + stream); + TORCH_CHECK(status == cudaSuccess, + "BatchPrefillWithPagedKVCacheSM90Run failed with error: ", + cudaGetErrorString(status)); + return true; + }); + }); +} diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_0.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..db8cade58529f1f4915e3a3b2d6900719bcfdb44 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_0.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchFP8PrefillWithPagedKVCacheDispatched + <128, + MaskMode::kNone, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultFP8Attention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchFP8PrefillWithPagedKVCacheDispatched + <128, + MaskMode::kNone, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultFP8Attention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_1.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..c0982d059808c912ee27445f54fb4630011f9d46 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_1.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchFP8PrefillWithPagedKVCacheDispatched + <128, + MaskMode::kCausal, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultFP8Attention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchFP8PrefillWithPagedKVCacheDispatched + <128, + MaskMode::kCausal, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultFP8Attention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_2.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..1116f67965b40e78d68b33f975b141cbd98159fa --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_2.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchFP8PrefillWithPagedKVCacheDispatched + <128, + MaskMode::kCustom, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultFP8Attention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchFP8PrefillWithPagedKVCacheDispatched + <128, + MaskMode::kCustom, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultFP8Attention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_3.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..629cdbf5a1bdea4b6c380a6b4f150e58913b3548 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_3.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchFP8PrefillWithPagedKVCacheDispatched + <128, + MaskMode::kMultiItemScoring, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultFP8Attention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchFP8PrefillWithPagedKVCacheDispatched + <128, + MaskMode::kMultiItemScoring, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultFP8Attention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_0.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..730347d4abcaff413b5286e36cc0dfe09228751d --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_0.cu @@ -0,0 +1 @@ +// TODO: Not implemented yet \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_1.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..730347d4abcaff413b5286e36cc0dfe09228751d --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_1.cu @@ -0,0 +1 @@ +// TODO: Not implemented yet \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_2.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..730347d4abcaff413b5286e36cc0dfe09228751d --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_2.cu @@ -0,0 +1 @@ +// TODO: Not implemented yet \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_3.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..730347d4abcaff413b5286e36cc0dfe09228751d --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_3.cu @@ -0,0 +1 @@ +// TODO: Not implemented yet \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_sm90_config.inc b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_sm90_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..bee5f05b2b2951d4b295b86df6e6f6e46eb4fe4d --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_sm90_config.inc @@ -0,0 +1,129 @@ +#pragma once +#include +#include +#include +#include +#include +#include "pytorch_extension_utils.h" + +#define ADDITIONAL_FUNC_PARAMS , at::Tensor scale_q, at::Tensor scale_k, at::Tensor scale_v, double sm_scale +#define ADDITIONAL_PARAMS_SETTER params.additional_params.scale_q = static_cast(scale_q.data_ptr()); \ +params.additional_params.scale_k = static_cast(scale_k.data_ptr()); \ +params.additional_params.scale_v = static_cast(scale_v.data_ptr()); \ +params.additional_params.sm_scale = sm_scale; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, RaggedParams, PagedParams, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { using AttentionVariant = DefaultFP8Attention; __VA_ARGS__();}) + +using namespace flashinfer; + +using DTypeQ = cutlass_dtype_t<__nv_fp8_e4m3>; +using DTypeKV = cutlass_dtype_t<__nv_fp8_e4m3>; +using DTypeO = cutlass_dtype_t; +using IdType = cutlass_dtype_t; + +constexpr int HEAD_DIM_QK = 128; +constexpr int HEAD_DIM_VO = 128; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto USE_SLIDING_WINDOW = false; + +struct RaggedParams { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + // The QKV matrices. + DTypeQ* q_ptr; + DTypeKV* k_ptr; + DTypeKV* v_ptr; + DTypeO* o_ptr; + float* lse_ptr; + + IdType* qo_tile_indices; + IdType* qo_indptr; + IdType* kv_indptr; + IdType* qo_lens; + IdType* kv_lens; + IdType* head_indices; + IdType* work_indptr; + IdType* batch_indices; + + struct AdditionalParams { + float* scale_q; +float* scale_k; +float* scale_v; +double sm_scale; + + } additional_params; + + int64_t q_stride_n; + int64_t k_stride_n; + int64_t v_stride_n; + int64_t o_stride_n; + int64_t q_stride_h; + int64_t k_stride_h; + int64_t v_stride_h; + int64_t o_stride_h; + int64_t nnz_qo; + int64_t nnz_kv; + + int head_dim; + int num_qo_heads; + int num_kv_heads; + int group_size; + int window_left; + + bool causal; +}; + +struct PagedParams { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + // The QKV matrices. + DTypeQ* q_ptr; + DTypeKV* k_ptr; + DTypeKV* v_ptr; + DTypeO* o_ptr; + float* lse_ptr; + + IdType* qo_tile_indices; + IdType* qo_indptr; + IdType* kv_indptr; + IdType* kv_indices; + IdType* qo_lens; + IdType* kv_lens; + IdType* head_indices; + IdType* work_indptr; + IdType* batch_indices; + + struct AdditionalParams { + float* scale_q; +float* scale_k; +float* scale_v; +double sm_scale; + + } additional_params; + + int64_t q_stride_n; + int64_t k_stride_n; + int64_t v_stride_n; + int64_t o_stride_n; + int64_t q_stride_h; + int64_t k_stride_h; + int64_t v_stride_h; + int64_t o_stride_h; + int64_t nnz_qo; + + int head_dim; + int num_qo_heads; + int num_kv_heads; + int group_size; + int page_size; + int window_left; + + bool causal; +}; + +#include \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_sm90_jit_pybind.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_sm90_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..b1edd582549e124c776f2cf9b4503f4bf60d16ae --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_sm90_jit_pybind.cu @@ -0,0 +1,46 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "batch_prefill_sm90_config.inc" +#include "pytorch_extension_utils.h" + +at::Tensor BatchPrefillWithKVCacheSM90Plan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, + int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, + int64_t head_dim_vo, bool causal); + +void BatchPrefillWithRaggedKVCacheSM90Run( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +void BatchPrefillWithPagedKVCacheSM90Run( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, + at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Batch-request prefill attention with KV-Cache plan + m.def("plan", BatchPrefillWithKVCacheSM90Plan); + // Batch-request prefill attention with KV-Cache operator + m.def("ragged_run", BatchPrefillWithRaggedKVCacheSM90Run); + // Batch-request prefill attention with KV-Cache operator + m.def("paged_run", BatchPrefillWithPagedKVCacheSM90Run); +} diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_fp8_sm90.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_fp8_sm90.cu new file mode 100644 index 0000000000000000000000000000000000000000..9f356754a8108a2e095d971c14a93ef79f7cfb4a --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_fp8_sm90.cu @@ -0,0 +1,185 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include +#include + +#include "batch_prefill_sm90_config.inc" +#include "pytorch_conversion_utils.h" +#include "pytorch_extension_utils.h" + +namespace flashinfer { + +template +cudaError_t BatchFP8PrefillWithPagedKVCacheDispatched(Params& params, bool enable_pdl, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +at::Tensor BatchPrefillWithKVCacheSM90Plan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, + int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, + int64_t head_dim_vo, bool causal) { + size_t float_workspace_size_in_bytes = + float_workspace_buffer.size(0) * float_workspace_buffer.element_size(); + size_t int_workspace_size_in_bytes = + int_workspace_buffer.size(0) * int_workspace_buffer.element_size(); + + flashinfer::PrefillPlanSM90Info plan_info; + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + cudaError_t status = + PrefillSM90Plan(float_workspace_buffer.data_ptr(), float_workspace_size_in_bytes, + int_workspace_buffer.data_ptr(), page_locked_int_workspace_buffer.data_ptr(), + int_workspace_size_in_bytes, plan_info, qo_indptr.data_ptr(), + kv_indptr.data_ptr(), kv_len_arr.data_ptr(), total_num_rows, + batch_size, num_qo_heads, num_kv_heads, head_dim_qk, head_dim_vo, page_size, + causal, enable_cuda_graph, /*sizeof_dtype_o=*/2, stream); + + TORCH_CHECK(status == cudaSuccess, + "PrefillSM90Plan failed with error: ", cudaGetErrorString(status)); + + return vec_to_tensor(plan_info.ToVector()); +} + +void BatchPrefillWithRaggedKVCacheSM90Run(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor k, at::Tensor v, + at::Tensor qo_indptr, at::Tensor kv_indptr, at::Tensor o, + std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left, + bool enable_pdl // placeholder + ADDITIONAL_FUNC_PARAMS) { + return; // TODO: Implement this function +} + +void BatchPrefillWithPagedKVCacheSM90Run( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, + at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + PrefillPlanSM90Info plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == q.size(0), lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == q.size(1), lse.size(1), q.size(1)); + } + QKVLayout kv_layout = static_cast(layout); + int64_t num_kv_heads, page_size; + int64_t head_dim_qk = q.size(2); + int64_t head_dim_vo = paged_v_cache.size(3); + if (kv_layout == QKVLayout::kHND) { + num_kv_heads = paged_k_cache.size(1); + page_size = paged_k_cache.size(2); + } else { + page_size = paged_k_cache.size(1); + num_kv_heads = paged_k_cache.size(2); + } + + void* float_buffer_ptr = float_workspace_buffer.data_ptr(); + void* int_buffer_ptr = int_workspace_buffer.data_ptr(); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = paged_k_cache.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + const MaskMode mask_mode = static_cast(mask_mode_code); + bool use_swa = window_left != -1; + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, USE_SLIDING_WINDOW, + USE_LOGITS_SOFT_CAP, AttentionVariant, RaggedParams, PagedParams, [&] { + PagedParams params; + + params.q_ptr = static_cast(q.data_ptr()); + params.k_ptr = static_cast(paged_k_cache.data_ptr()); + params.v_ptr = static_cast(paged_v_cache.data_ptr()); + params.o_ptr = static_cast(o.data_ptr()); + params.lse_ptr = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.q_stride_n = q.stride(0); + params.q_stride_h = q.stride(1); + params.o_stride_n = o.stride(0); + params.o_stride_h = o.stride(1); + if (kv_layout == QKVLayout::kNHD) { + // (num_pages, page_size, num_heads, head_dim) + params.k_stride_n = paged_k_cache.stride(1); + params.k_stride_h = paged_k_cache.stride(2); + params.v_stride_n = paged_v_cache.stride(1); + params.v_stride_h = paged_v_cache.stride(2); + } else { + // (num_pages, num_heads, page_size, head_dim) + params.k_stride_h = paged_k_cache.stride(1); + params.k_stride_n = paged_k_cache.stride(2); + params.v_stride_h = paged_v_cache.stride(1); + params.v_stride_n = paged_v_cache.stride(2); + } + params.nnz_qo = q.size(0); + params.num_qo_heads = q.size(1); + params.num_kv_heads = num_kv_heads; + params.group_size = params.num_qo_heads / num_kv_heads; + params.page_size = page_size; + params.window_left = window_left; + params.causal = mask_mode_code == 1; + params.qo_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_tile_indices_offset); + params.qo_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_indptr_offset); + params.kv_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_indptr_offset); + params.qo_lens = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_len_offset); + params.kv_lens = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_len_offset); + params.batch_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.batch_indices_offset); + params.head_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.head_indices_offset); + params.work_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.work_indptr_offset); + params.kv_indices = static_cast(paged_kv_indices.data_ptr()); + + ADDITIONAL_PARAMS_SETTER + + // Not support various head_dim for now + static_assert(HEAD_DIM_QK == HEAD_DIM_VO, "head_dim_qk and head_dim_vo should be the same"); + // Currently only support same quantization precision + static_assert(std::is_same_v); + + bool same_schedule_for_all_heads = plan_info.same_schedule_for_all_heads; + DISPATCH_BOOL(same_schedule_for_all_heads, SAME_SCHEDULER_FOR_ALL_HEADS, [&] { + cudaError_t status = + BatchFP8PrefillWithPagedKVCacheDispatched(params, enable_pdl, + stream); + TORCH_CHECK(status == cudaSuccess, + "BatchPrefillWithPagedKVCacheSM90Run failed with error: ", + cudaGetErrorString(status)); + return true; + }); + }); +} diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_0.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..12fe96fede547f29a056b9925ee32001e82f49fd --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_0.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchFP8PrefillWithPagedKVCacheDispatched + <256, + MaskMode::kNone, + /*USE_SLIDING_WINDOW=*/true, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultFP8Attention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchFP8PrefillWithPagedKVCacheDispatched + <256, + MaskMode::kNone, + /*USE_SLIDING_WINDOW=*/true, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultFP8Attention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_1.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..53de9f1d9b0288a3a35245f794c09ab084a06271 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_1.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchFP8PrefillWithPagedKVCacheDispatched + <256, + MaskMode::kCausal, + /*USE_SLIDING_WINDOW=*/true, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultFP8Attention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchFP8PrefillWithPagedKVCacheDispatched + <256, + MaskMode::kCausal, + /*USE_SLIDING_WINDOW=*/true, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultFP8Attention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_2.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..d3854bf6989d9916f5045f1cc77b0504547219fa --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_2.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchFP8PrefillWithPagedKVCacheDispatched + <256, + MaskMode::kCustom, + /*USE_SLIDING_WINDOW=*/true, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultFP8Attention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchFP8PrefillWithPagedKVCacheDispatched + <256, + MaskMode::kCustom, + /*USE_SLIDING_WINDOW=*/true, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultFP8Attention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_3.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..fac46d2f0439c7c36c9b95169b3b864a86c7bb13 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_3.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchFP8PrefillWithPagedKVCacheDispatched + <256, + MaskMode::kMultiItemScoring, + /*USE_SLIDING_WINDOW=*/true, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultFP8Attention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchFP8PrefillWithPagedKVCacheDispatched + <256, + MaskMode::kMultiItemScoring, + /*USE_SLIDING_WINDOW=*/true, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultFP8Attention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_0.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..730347d4abcaff413b5286e36cc0dfe09228751d --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_0.cu @@ -0,0 +1 @@ +// TODO: Not implemented yet \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_1.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..730347d4abcaff413b5286e36cc0dfe09228751d --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_1.cu @@ -0,0 +1 @@ +// TODO: Not implemented yet \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_2.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..730347d4abcaff413b5286e36cc0dfe09228751d --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_2.cu @@ -0,0 +1 @@ +// TODO: Not implemented yet \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_3.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..730347d4abcaff413b5286e36cc0dfe09228751d --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_3.cu @@ -0,0 +1 @@ +// TODO: Not implemented yet \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_sm90_config.inc b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_sm90_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..bf9775b33d1f72a35d0f1b3c4366e71c16ba12e3 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_sm90_config.inc @@ -0,0 +1,129 @@ +#pragma once +#include +#include +#include +#include +#include +#include "pytorch_extension_utils.h" + +#define ADDITIONAL_FUNC_PARAMS , at::Tensor scale_q, at::Tensor scale_k, at::Tensor scale_v, double sm_scale +#define ADDITIONAL_PARAMS_SETTER params.additional_params.scale_q = static_cast(scale_q.data_ptr()); \ +params.additional_params.scale_k = static_cast(scale_k.data_ptr()); \ +params.additional_params.scale_v = static_cast(scale_v.data_ptr()); \ +params.additional_params.sm_scale = sm_scale; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, RaggedParams, PagedParams, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { using AttentionVariant = DefaultFP8Attention; __VA_ARGS__();}) + +using namespace flashinfer; + +using DTypeQ = cutlass_dtype_t<__nv_fp8_e4m3>; +using DTypeKV = cutlass_dtype_t<__nv_fp8_e4m3>; +using DTypeO = cutlass_dtype_t; +using IdType = cutlass_dtype_t; + +constexpr int HEAD_DIM_QK = 256; +constexpr int HEAD_DIM_VO = 256; +constexpr auto USE_LOGITS_SOFT_CAP = true; +constexpr auto USE_SLIDING_WINDOW = true; + +struct RaggedParams { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + // The QKV matrices. + DTypeQ* q_ptr; + DTypeKV* k_ptr; + DTypeKV* v_ptr; + DTypeO* o_ptr; + float* lse_ptr; + + IdType* qo_tile_indices; + IdType* qo_indptr; + IdType* kv_indptr; + IdType* qo_lens; + IdType* kv_lens; + IdType* head_indices; + IdType* work_indptr; + IdType* batch_indices; + + struct AdditionalParams { + float* scale_q; +float* scale_k; +float* scale_v; +double sm_scale; + + } additional_params; + + int64_t q_stride_n; + int64_t k_stride_n; + int64_t v_stride_n; + int64_t o_stride_n; + int64_t q_stride_h; + int64_t k_stride_h; + int64_t v_stride_h; + int64_t o_stride_h; + int64_t nnz_qo; + int64_t nnz_kv; + + int head_dim; + int num_qo_heads; + int num_kv_heads; + int group_size; + int window_left; + + bool causal; +}; + +struct PagedParams { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + // The QKV matrices. + DTypeQ* q_ptr; + DTypeKV* k_ptr; + DTypeKV* v_ptr; + DTypeO* o_ptr; + float* lse_ptr; + + IdType* qo_tile_indices; + IdType* qo_indptr; + IdType* kv_indptr; + IdType* kv_indices; + IdType* qo_lens; + IdType* kv_lens; + IdType* head_indices; + IdType* work_indptr; + IdType* batch_indices; + + struct AdditionalParams { + float* scale_q; +float* scale_k; +float* scale_v; +double sm_scale; + + } additional_params; + + int64_t q_stride_n; + int64_t k_stride_n; + int64_t v_stride_n; + int64_t o_stride_n; + int64_t q_stride_h; + int64_t k_stride_h; + int64_t v_stride_h; + int64_t o_stride_h; + int64_t nnz_qo; + + int head_dim; + int num_qo_heads; + int num_kv_heads; + int group_size; + int page_size; + int window_left; + + bool causal; +}; + +#include \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_sm90_jit_pybind.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_sm90_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..b1edd582549e124c776f2cf9b4503f4bf60d16ae --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_sm90_jit_pybind.cu @@ -0,0 +1,46 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "batch_prefill_sm90_config.inc" +#include "pytorch_extension_utils.h" + +at::Tensor BatchPrefillWithKVCacheSM90Plan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, + int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, + int64_t head_dim_vo, bool causal); + +void BatchPrefillWithRaggedKVCacheSM90Run( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +void BatchPrefillWithPagedKVCacheSM90Run( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, + at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Batch-request prefill attention with KV-Cache plan + m.def("plan", BatchPrefillWithKVCacheSM90Plan); + // Batch-request prefill attention with KV-Cache operator + m.def("ragged_run", BatchPrefillWithRaggedKVCacheSM90Run); + // Batch-request prefill attention with KV-Cache operator + m.def("paged_run", BatchPrefillWithPagedKVCacheSM90Run); +} diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill.cu new file mode 100644 index 0000000000000000000000000000000000000000..a51fc7f56a0af7482bd5d2a9d3f672854df69dfd --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill.cu @@ -0,0 +1,330 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include +#include + +#include "batch_prefill_config.inc" +#include "pytorch_conversion_utils.h" +#include "pytorch_extension_utils.h" + +namespace flashinfer { + +template +cudaError_t BatchPrefillWithPagedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v, + float* tmp_s, bool enable_pdl, + cudaStream_t stream); + +template +cudaError_t BatchPrefillWithRaggedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v, + float* tmp_s, bool enable_pdl, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +at::Tensor BatchPrefillWithKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, + int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, + int64_t head_dim_vo, bool causal) { + size_t float_workspace_size_in_bytes = + float_workspace_buffer.size(0) * float_workspace_buffer.element_size(); + size_t int_workspace_size_in_bytes = + int_workspace_buffer.size(0) * int_workspace_buffer.element_size(); + + PrefillPlanInfo plan_info; + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + cudaError_t status = PrefillPlan( + float_workspace_buffer.data_ptr(), float_workspace_size_in_bytes, + int_workspace_buffer.data_ptr(), page_locked_int_workspace_buffer.data_ptr(), + int_workspace_size_in_bytes, plan_info, qo_indptr.data_ptr(), + kv_indptr.data_ptr(), total_num_rows, batch_size, num_qo_heads, num_kv_heads, + head_dim_qk, head_dim_vo, page_size, enable_cuda_graph, /*sizeof_dtype_o=*/2, stream); + + TORCH_CHECK(status == cudaSuccess, + "Failed to plan prefill with error: ", cudaGetErrorString(status)); + + return vec_to_tensor(plan_info.ToVector()); +} + +void BatchPrefillWithRaggedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor k, at::Tensor v, + at::Tensor qo_indptr, at::Tensor kv_indptr, at::Tensor o, + std::optional maybe_lse, int64_t mask_mode_code, + int64_t layout, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + PrefillPlanInfo plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + QKVLayout kv_layout = static_cast(layout); + + int64_t num_qo_heads = q.size(1); + int64_t head_dim_qk = q.size(2); + int64_t num_kv_heads = (kv_layout == QKVLayout::kNHD) ? k.size(1) : k.size(0); + uint32_t q_stride_n = q.stride(0), q_stride_h = q.stride(1), k_stride_n, k_stride_h, v_stride_n, + v_stride_h; + if (kv_layout == QKVLayout::kNHD) { + k_stride_n = k.stride(0); + k_stride_h = k.stride(1); + v_stride_n = v.stride(0); + v_stride_h = v.stride(1); + } else { + k_stride_h = k.stride(0); + k_stride_n = k.stride(1); + v_stride_h = v.stride(0); + v_stride_n = v.stride(1); + } + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == q.size(0), lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == q.size(1), lse.size(1), q.size(1)); + } + + void* float_buffer_ptr = float_workspace_buffer.data_ptr(); + void* int_buffer_ptr = int_workspace_buffer.data_ptr(); + + const MaskMode mask_mode = static_cast(mask_mode_code); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, + RaggedParams, PagedParams, [&] { + RaggedParams params; + + params.q = static_cast(q.data_ptr()); + params.k = static_cast(k.data_ptr()); + params.v = static_cast(v.data_ptr()); + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.q_indptr = static_cast(qo_indptr.data_ptr()); + params.kv_indptr = static_cast(kv_indptr.data_ptr()); + params.num_qo_heads = num_qo_heads; + params.num_kv_heads = num_kv_heads; + params.group_size = uint_fastdiv(num_qo_heads / num_kv_heads); + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.k_stride_n = k_stride_n; + params.k_stride_h = k_stride_h; + params.v_stride_n = v_stride_n; + params.v_stride_h = v_stride_h; + params.window_left = window_left; + + params.request_indices = nullptr; + params.qo_tile_indices = nullptr; + params.kv_tile_indices = nullptr; + params.merge_indptr = nullptr; + params.o_indptr = nullptr; + params.kv_chunk_size_ptr = nullptr; + params.block_valid_mask = nullptr; + params.total_num_rows = nullptr; + params.max_total_num_rows = 0; + params.padded_batch_size = 0; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + DTypeO* tmp_v = nullptr; + float* tmp_s = nullptr; + + params.request_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.request_indices_offset); + params.qo_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_tile_indices_offset); + params.kv_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_tile_indices_offset); + params.o_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.o_indptr_offset); + params.kv_chunk_size_ptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_chunk_size_ptr_offset); + if (plan_info.split_kv) { + params.merge_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.merge_indptr_offset); + tmp_v = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.v_offset); + tmp_s = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.s_offset); + if (plan_info.enable_cuda_graph) { + params.block_valid_mask = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.block_valid_mask_offset); + } + } + params.padded_batch_size = plan_info.padded_batch_size; + params.max_total_num_rows = plan_info.total_num_rows; + if (plan_info.enable_cuda_graph) { + params.total_num_rows = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.total_num_rows_offset); + } + + cudaError_t status = cudaSuccess; + + DISPATCH_CTA_TILE_Q(plan_info.cta_tile_q, CTA_TILE_Q, { + status = flashinfer::BatchPrefillWithRaggedKVCacheDispatched< + CTA_TILE_Q, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + /*use_fp16_qk_reduction=*/USE_FP16_QK_REDUCTION, MASK_MODE, AttentionVariant, + RaggedParams>(params, tmp_v, tmp_s, enable_pdl, stream); + }); + + TORCH_CHECK(status == cudaSuccess, "BatchPrefillWithRaggedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); +} + +void BatchPrefillWithPagedKVCacheRun( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, + at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + PrefillPlanInfo plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + QKVLayout kv_layout = static_cast(layout); + auto device = q.device(); + int64_t batch_size = paged_kv_indptr.size(0) - 1; + int64_t num_qo_heads = q.size(1); + int64_t num_kv_heads, page_size; + uint32_t head_dim_qk = q.size(2); + if (kv_layout == QKVLayout::kHND) { + num_kv_heads = paged_k_cache.size(1); + page_size = paged_k_cache.size(2); + } else { + page_size = paged_k_cache.size(1); + num_kv_heads = paged_k_cache.size(2); + } + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == q.size(0), lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == q.size(1), lse.size(1), q.size(1)); + } + + void* float_buffer_ptr = static_cast(float_workspace_buffer.data_ptr()); + void* int_buffer_ptr = static_cast(int_workspace_buffer.data_ptr()); + + const MaskMode mask_mode = static_cast(mask_mode_code); + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = paged_k_cache.scalar_type(); + + // get q_stride_n and q_stride_h + const auto q_stride_n = q.stride(0); + const auto q_stride_h = q.stride(1); + + // get kv_cache_strides + const int64_t* kv_cache_strides = nullptr; + auto k_strides = paged_k_cache.strides(); + auto v_strides = paged_v_cache.strides(); + TORCH_CHECK(k_strides == v_strides, "k/v strides must be identical"); + kv_cache_strides = k_strides.data(); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, + RaggedParams, PagedParams, [&] { + PagedParams params; + + params.q = static_cast(q.data_ptr()); + paged_kv_t paged_kv( + num_kv_heads, page_size, HEAD_DIM_VO, batch_size, kv_layout, + static_cast(paged_k_cache.data_ptr()), + static_cast(paged_v_cache.data_ptr()), kv_cache_strides, + static_cast(paged_kv_indices.data_ptr()), + static_cast(paged_kv_indptr.data_ptr()), + static_cast(paged_kv_last_page_len.data_ptr())); + params.paged_kv = paged_kv; + params.q_indptr = static_cast(qo_indptr.data_ptr()); + params.o = static_cast(o.data_ptr()); + + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.num_qo_heads = num_qo_heads; + params.group_size = uint_fastdiv(num_qo_heads / paged_kv.num_heads); + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.window_left = window_left; + + params.request_indices = nullptr; + params.qo_tile_indices = nullptr; + params.kv_tile_indices = nullptr; + params.merge_indptr = nullptr; + params.o_indptr = nullptr; + params.kv_chunk_size_ptr = nullptr; + params.block_valid_mask = nullptr; + params.total_num_rows = nullptr; + params.max_total_num_rows = 0; + params.padded_batch_size = 0; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + DTypeO* tmp_v = nullptr; + float* tmp_s = nullptr; + + params.request_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.request_indices_offset); + params.qo_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_tile_indices_offset); + params.kv_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_tile_indices_offset); + params.o_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.o_indptr_offset); + params.kv_chunk_size_ptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_chunk_size_ptr_offset); + if (plan_info.split_kv) { + params.merge_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.merge_indptr_offset); + tmp_v = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.v_offset); + tmp_s = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.s_offset); + if (plan_info.enable_cuda_graph) { + params.block_valid_mask = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.block_valid_mask_offset); + } + } + params.padded_batch_size = plan_info.padded_batch_size; + params.max_total_num_rows = plan_info.total_num_rows; + if (plan_info.enable_cuda_graph) { + params.total_num_rows = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.total_num_rows_offset); + } + + cudaError_t status = cudaSuccess; + + DISPATCH_CTA_TILE_Q(plan_info.cta_tile_q, CTA_TILE_Q, { + status = flashinfer::BatchPrefillWithPagedKVCacheDispatched< + CTA_TILE_Q, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + /*use_fp16_qk_reduction=*/USE_FP16_QK_REDUCTION, MASK_MODE, AttentionVariant, + PagedParams>(params, tmp_v, tmp_s, enable_pdl, stream); + }); + + TORCH_CHECK(status == cudaSuccess, "BatchPrefillWithPagedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); +} diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_config.inc b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..59ddaee0af14816c5585520c1612f07e5fdaed2c --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_config.inc @@ -0,0 +1,153 @@ +#pragma once +#include +#include +#include +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_custom_mask, std::optional maybe_mask_indptr, std::optional maybe_alibi_slopes, std::optional maybe_prefix_len_ptr, std::optional maybe_token_pos_in_items_ptr, std::optional maybe_max_item_len_ptr, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta, int64_t token_pos_in_items_len +#define ADDITIONAL_PARAMS_SETTER params.maybe_custom_mask = maybe_custom_mask ? static_cast(maybe_custom_mask->data_ptr()): nullptr; \ +params.maybe_mask_indptr = maybe_mask_indptr ? static_cast(maybe_mask_indptr->data_ptr()): nullptr; \ +params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.maybe_prefix_len_ptr = maybe_prefix_len_ptr ? static_cast(maybe_prefix_len_ptr->data_ptr()): nullptr; \ +params.maybe_token_pos_in_items_ptr = maybe_token_pos_in_items_ptr ? static_cast(maybe_token_pos_in_items_ptr->data_ptr()): nullptr; \ +params.maybe_max_item_len_ptr = maybe_max_item_len_ptr ? static_cast(maybe_max_item_len_ptr->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; \ +params.token_pos_in_items_len = token_pos_in_items_len; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, RaggedParams, PagedParams, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { \ + constexpr auto use_custom_mask = MASK_MODE == MaskMode::kCustom; \ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ + }) + +using namespace flashinfer; + +using DTypeQ = half; +using DTypeKV = __nv_fp8_e4m3; +using DTypeO = half; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 128; +constexpr int HEAD_DIM_VO = 128; +constexpr bool USE_FP16_QK_REDUCTION = false; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = false; + + +struct RaggedParams { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + DTypeQ* q; + DTypeKV* k; + DTypeKV* v; + IdType* q_indptr; + IdType* kv_indptr; + DTypeO* o; + float* lse; + uint_fastdiv group_size; + + uint8_t* maybe_custom_mask; +int32_t* maybe_mask_indptr; +float* maybe_alibi_slopes; +uint32_t* maybe_prefix_len_ptr; +uint16_t* maybe_token_pos_in_items_ptr; +uint16_t* maybe_max_item_len_ptr; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; +int64_t token_pos_in_items_len; + + uint32_t num_qo_heads; + uint32_t num_kv_heads; + uint32_t q_stride_n; + uint32_t q_stride_h; + uint32_t k_stride_n; + uint32_t k_stride_h; + uint32_t v_stride_n; + uint32_t v_stride_h; + int32_t window_left; + + IdType* request_indices; + IdType* qo_tile_indices; + IdType* kv_tile_indices; + IdType* merge_indptr; + IdType* o_indptr; + IdType* kv_chunk_size_ptr; + bool* block_valid_mask; + uint32_t max_total_num_rows; + uint32_t* total_num_rows; + uint32_t padded_batch_size; + bool partition_kv; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { + return q_indptr[batch_idx + 1] - q_indptr[batch_idx]; + } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return kv_indptr[batch_idx + 1] - kv_indptr[batch_idx]; + } +}; + +struct PagedParams { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + DTypeQ* q; + paged_kv_t paged_kv; + IdType* q_indptr; + DTypeO* o; + float* lse; + uint_fastdiv group_size; + + uint8_t* maybe_custom_mask; +int32_t* maybe_mask_indptr; +float* maybe_alibi_slopes; +uint32_t* maybe_prefix_len_ptr; +uint16_t* maybe_token_pos_in_items_ptr; +uint16_t* maybe_max_item_len_ptr; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; +int64_t token_pos_in_items_len; + + uint32_t num_qo_heads; + IdType q_stride_n; + IdType q_stride_h; + int32_t window_left; + + IdType* request_indices; + IdType* qo_tile_indices; + IdType* kv_tile_indices; + IdType* merge_indptr; + IdType* o_indptr; + bool* block_valid_mask; + IdType* kv_chunk_size_ptr; + uint32_t max_total_num_rows; + uint32_t* total_num_rows; + uint32_t padded_batch_size; + bool partition_kv; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { + return q_indptr[batch_idx + 1] - q_indptr[batch_idx]; + } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return paged_kv.get_length(batch_idx); + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_jit_pybind.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..5421ab1cf910fdfbc867c182c8e22f23e44749ae --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_jit_pybind.cu @@ -0,0 +1,48 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "batch_prefill_config.inc" +#include "pytorch_extension_utils.h" + +at::Tensor BatchPrefillWithKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, + int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, + int64_t head_dim_vo, bool causal); + +void BatchPrefillWithRaggedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor k, at::Tensor v, + at::Tensor qo_indptr, at::Tensor kv_indptr, at::Tensor o, + std::optional maybe_lse, int64_t mask_mode_code, + int64_t layout, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +void BatchPrefillWithPagedKVCacheRun( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, + at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Batch-request prefill attention with KV-Cache plan + m.def("plan", BatchPrefillWithKVCachePlan); + // Batch-request prefill attention with KV-Cache operator + m.def("ragged_run", BatchPrefillWithRaggedKVCacheRun); + // Batch-request prefill attention with KV-Cache operator + m.def("paged_run", BatchPrefillWithPagedKVCacheRun); +} diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_0.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..ff4fbd96e95fa670c019e4fb5cda473f41ba1628 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_0.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kNone == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 128, 128, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 128, 128, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 128, 128, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_1.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..ae257f99a6b4f7b44ff90ee4fb0395a49dedaf9c --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_1.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCausal == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 128, 128, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 128, 128, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 128, 128, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_2.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..9baac31703e1a9977eb8e85ef79bd23e4634ae5c --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_2.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCustom == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 128, 128, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 128, 128, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 128, 128, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_3.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..108509bf3bd71e251f388881bca3f7f7f26154e7 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_3.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kMultiItemScoring == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 128, 128, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 128, 128, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 128, 128, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_0.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..39450198451b3c42f30ed9e3eb829e0be30fe4a1 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_0.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kNone == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 128, 128, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 128, 128, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 128, 128, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_1.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..50ab7b2f475eedf92137e309509db05f309cc591 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_1.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCausal == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 128, 128, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 128, 128, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 128, 128, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_2.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..86408ffce4c329caabf1f856d054120815065520 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_2.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCustom == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 128, 128, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 128, 128, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 128, 128, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_3.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..b44ab24580d5632515cf18460b859352aa95754e --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_3.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kMultiItemScoring == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 128, 128, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 128, 128, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 128, 128, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill.cu new file mode 100644 index 0000000000000000000000000000000000000000..a51fc7f56a0af7482bd5d2a9d3f672854df69dfd --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill.cu @@ -0,0 +1,330 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include +#include + +#include "batch_prefill_config.inc" +#include "pytorch_conversion_utils.h" +#include "pytorch_extension_utils.h" + +namespace flashinfer { + +template +cudaError_t BatchPrefillWithPagedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v, + float* tmp_s, bool enable_pdl, + cudaStream_t stream); + +template +cudaError_t BatchPrefillWithRaggedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v, + float* tmp_s, bool enable_pdl, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +at::Tensor BatchPrefillWithKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, + int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, + int64_t head_dim_vo, bool causal) { + size_t float_workspace_size_in_bytes = + float_workspace_buffer.size(0) * float_workspace_buffer.element_size(); + size_t int_workspace_size_in_bytes = + int_workspace_buffer.size(0) * int_workspace_buffer.element_size(); + + PrefillPlanInfo plan_info; + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + cudaError_t status = PrefillPlan( + float_workspace_buffer.data_ptr(), float_workspace_size_in_bytes, + int_workspace_buffer.data_ptr(), page_locked_int_workspace_buffer.data_ptr(), + int_workspace_size_in_bytes, plan_info, qo_indptr.data_ptr(), + kv_indptr.data_ptr(), total_num_rows, batch_size, num_qo_heads, num_kv_heads, + head_dim_qk, head_dim_vo, page_size, enable_cuda_graph, /*sizeof_dtype_o=*/2, stream); + + TORCH_CHECK(status == cudaSuccess, + "Failed to plan prefill with error: ", cudaGetErrorString(status)); + + return vec_to_tensor(plan_info.ToVector()); +} + +void BatchPrefillWithRaggedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor k, at::Tensor v, + at::Tensor qo_indptr, at::Tensor kv_indptr, at::Tensor o, + std::optional maybe_lse, int64_t mask_mode_code, + int64_t layout, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + PrefillPlanInfo plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + QKVLayout kv_layout = static_cast(layout); + + int64_t num_qo_heads = q.size(1); + int64_t head_dim_qk = q.size(2); + int64_t num_kv_heads = (kv_layout == QKVLayout::kNHD) ? k.size(1) : k.size(0); + uint32_t q_stride_n = q.stride(0), q_stride_h = q.stride(1), k_stride_n, k_stride_h, v_stride_n, + v_stride_h; + if (kv_layout == QKVLayout::kNHD) { + k_stride_n = k.stride(0); + k_stride_h = k.stride(1); + v_stride_n = v.stride(0); + v_stride_h = v.stride(1); + } else { + k_stride_h = k.stride(0); + k_stride_n = k.stride(1); + v_stride_h = v.stride(0); + v_stride_n = v.stride(1); + } + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == q.size(0), lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == q.size(1), lse.size(1), q.size(1)); + } + + void* float_buffer_ptr = float_workspace_buffer.data_ptr(); + void* int_buffer_ptr = int_workspace_buffer.data_ptr(); + + const MaskMode mask_mode = static_cast(mask_mode_code); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, + RaggedParams, PagedParams, [&] { + RaggedParams params; + + params.q = static_cast(q.data_ptr()); + params.k = static_cast(k.data_ptr()); + params.v = static_cast(v.data_ptr()); + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.q_indptr = static_cast(qo_indptr.data_ptr()); + params.kv_indptr = static_cast(kv_indptr.data_ptr()); + params.num_qo_heads = num_qo_heads; + params.num_kv_heads = num_kv_heads; + params.group_size = uint_fastdiv(num_qo_heads / num_kv_heads); + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.k_stride_n = k_stride_n; + params.k_stride_h = k_stride_h; + params.v_stride_n = v_stride_n; + params.v_stride_h = v_stride_h; + params.window_left = window_left; + + params.request_indices = nullptr; + params.qo_tile_indices = nullptr; + params.kv_tile_indices = nullptr; + params.merge_indptr = nullptr; + params.o_indptr = nullptr; + params.kv_chunk_size_ptr = nullptr; + params.block_valid_mask = nullptr; + params.total_num_rows = nullptr; + params.max_total_num_rows = 0; + params.padded_batch_size = 0; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + DTypeO* tmp_v = nullptr; + float* tmp_s = nullptr; + + params.request_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.request_indices_offset); + params.qo_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_tile_indices_offset); + params.kv_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_tile_indices_offset); + params.o_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.o_indptr_offset); + params.kv_chunk_size_ptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_chunk_size_ptr_offset); + if (plan_info.split_kv) { + params.merge_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.merge_indptr_offset); + tmp_v = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.v_offset); + tmp_s = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.s_offset); + if (plan_info.enable_cuda_graph) { + params.block_valid_mask = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.block_valid_mask_offset); + } + } + params.padded_batch_size = plan_info.padded_batch_size; + params.max_total_num_rows = plan_info.total_num_rows; + if (plan_info.enable_cuda_graph) { + params.total_num_rows = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.total_num_rows_offset); + } + + cudaError_t status = cudaSuccess; + + DISPATCH_CTA_TILE_Q(plan_info.cta_tile_q, CTA_TILE_Q, { + status = flashinfer::BatchPrefillWithRaggedKVCacheDispatched< + CTA_TILE_Q, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + /*use_fp16_qk_reduction=*/USE_FP16_QK_REDUCTION, MASK_MODE, AttentionVariant, + RaggedParams>(params, tmp_v, tmp_s, enable_pdl, stream); + }); + + TORCH_CHECK(status == cudaSuccess, "BatchPrefillWithRaggedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); +} + +void BatchPrefillWithPagedKVCacheRun( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, + at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + PrefillPlanInfo plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + QKVLayout kv_layout = static_cast(layout); + auto device = q.device(); + int64_t batch_size = paged_kv_indptr.size(0) - 1; + int64_t num_qo_heads = q.size(1); + int64_t num_kv_heads, page_size; + uint32_t head_dim_qk = q.size(2); + if (kv_layout == QKVLayout::kHND) { + num_kv_heads = paged_k_cache.size(1); + page_size = paged_k_cache.size(2); + } else { + page_size = paged_k_cache.size(1); + num_kv_heads = paged_k_cache.size(2); + } + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == q.size(0), lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == q.size(1), lse.size(1), q.size(1)); + } + + void* float_buffer_ptr = static_cast(float_workspace_buffer.data_ptr()); + void* int_buffer_ptr = static_cast(int_workspace_buffer.data_ptr()); + + const MaskMode mask_mode = static_cast(mask_mode_code); + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = paged_k_cache.scalar_type(); + + // get q_stride_n and q_stride_h + const auto q_stride_n = q.stride(0); + const auto q_stride_h = q.stride(1); + + // get kv_cache_strides + const int64_t* kv_cache_strides = nullptr; + auto k_strides = paged_k_cache.strides(); + auto v_strides = paged_v_cache.strides(); + TORCH_CHECK(k_strides == v_strides, "k/v strides must be identical"); + kv_cache_strides = k_strides.data(); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, + RaggedParams, PagedParams, [&] { + PagedParams params; + + params.q = static_cast(q.data_ptr()); + paged_kv_t paged_kv( + num_kv_heads, page_size, HEAD_DIM_VO, batch_size, kv_layout, + static_cast(paged_k_cache.data_ptr()), + static_cast(paged_v_cache.data_ptr()), kv_cache_strides, + static_cast(paged_kv_indices.data_ptr()), + static_cast(paged_kv_indptr.data_ptr()), + static_cast(paged_kv_last_page_len.data_ptr())); + params.paged_kv = paged_kv; + params.q_indptr = static_cast(qo_indptr.data_ptr()); + params.o = static_cast(o.data_ptr()); + + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.num_qo_heads = num_qo_heads; + params.group_size = uint_fastdiv(num_qo_heads / paged_kv.num_heads); + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.window_left = window_left; + + params.request_indices = nullptr; + params.qo_tile_indices = nullptr; + params.kv_tile_indices = nullptr; + params.merge_indptr = nullptr; + params.o_indptr = nullptr; + params.kv_chunk_size_ptr = nullptr; + params.block_valid_mask = nullptr; + params.total_num_rows = nullptr; + params.max_total_num_rows = 0; + params.padded_batch_size = 0; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + DTypeO* tmp_v = nullptr; + float* tmp_s = nullptr; + + params.request_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.request_indices_offset); + params.qo_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_tile_indices_offset); + params.kv_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_tile_indices_offset); + params.o_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.o_indptr_offset); + params.kv_chunk_size_ptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_chunk_size_ptr_offset); + if (plan_info.split_kv) { + params.merge_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.merge_indptr_offset); + tmp_v = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.v_offset); + tmp_s = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.s_offset); + if (plan_info.enable_cuda_graph) { + params.block_valid_mask = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.block_valid_mask_offset); + } + } + params.padded_batch_size = plan_info.padded_batch_size; + params.max_total_num_rows = plan_info.total_num_rows; + if (plan_info.enable_cuda_graph) { + params.total_num_rows = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.total_num_rows_offset); + } + + cudaError_t status = cudaSuccess; + + DISPATCH_CTA_TILE_Q(plan_info.cta_tile_q, CTA_TILE_Q, { + status = flashinfer::BatchPrefillWithPagedKVCacheDispatched< + CTA_TILE_Q, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + /*use_fp16_qk_reduction=*/USE_FP16_QK_REDUCTION, MASK_MODE, AttentionVariant, + PagedParams>(params, tmp_v, tmp_s, enable_pdl, stream); + }); + + TORCH_CHECK(status == cudaSuccess, "BatchPrefillWithPagedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); +} diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_config.inc b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..bc4c9b0b5d16335ee2bb11ca895641c9f7518a09 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_config.inc @@ -0,0 +1,153 @@ +#pragma once +#include +#include +#include +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_custom_mask, std::optional maybe_mask_indptr, std::optional maybe_alibi_slopes, std::optional maybe_prefix_len_ptr, std::optional maybe_token_pos_in_items_ptr, std::optional maybe_max_item_len_ptr, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta, int64_t token_pos_in_items_len +#define ADDITIONAL_PARAMS_SETTER params.maybe_custom_mask = maybe_custom_mask ? static_cast(maybe_custom_mask->data_ptr()): nullptr; \ +params.maybe_mask_indptr = maybe_mask_indptr ? static_cast(maybe_mask_indptr->data_ptr()): nullptr; \ +params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.maybe_prefix_len_ptr = maybe_prefix_len_ptr ? static_cast(maybe_prefix_len_ptr->data_ptr()): nullptr; \ +params.maybe_token_pos_in_items_ptr = maybe_token_pos_in_items_ptr ? static_cast(maybe_token_pos_in_items_ptr->data_ptr()): nullptr; \ +params.maybe_max_item_len_ptr = maybe_max_item_len_ptr ? static_cast(maybe_max_item_len_ptr->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; \ +params.token_pos_in_items_len = token_pos_in_items_len; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, RaggedParams, PagedParams, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { \ + constexpr auto use_custom_mask = MASK_MODE == MaskMode::kCustom; \ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ + }) + +using namespace flashinfer; + +using DTypeQ = half; +using DTypeKV = __nv_fp8_e4m3; +using DTypeO = half; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 256; +constexpr int HEAD_DIM_VO = 256; +constexpr bool USE_FP16_QK_REDUCTION = false; +constexpr auto USE_LOGITS_SOFT_CAP = true; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = true; + + +struct RaggedParams { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + DTypeQ* q; + DTypeKV* k; + DTypeKV* v; + IdType* q_indptr; + IdType* kv_indptr; + DTypeO* o; + float* lse; + uint_fastdiv group_size; + + uint8_t* maybe_custom_mask; +int32_t* maybe_mask_indptr; +float* maybe_alibi_slopes; +uint32_t* maybe_prefix_len_ptr; +uint16_t* maybe_token_pos_in_items_ptr; +uint16_t* maybe_max_item_len_ptr; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; +int64_t token_pos_in_items_len; + + uint32_t num_qo_heads; + uint32_t num_kv_heads; + uint32_t q_stride_n; + uint32_t q_stride_h; + uint32_t k_stride_n; + uint32_t k_stride_h; + uint32_t v_stride_n; + uint32_t v_stride_h; + int32_t window_left; + + IdType* request_indices; + IdType* qo_tile_indices; + IdType* kv_tile_indices; + IdType* merge_indptr; + IdType* o_indptr; + IdType* kv_chunk_size_ptr; + bool* block_valid_mask; + uint32_t max_total_num_rows; + uint32_t* total_num_rows; + uint32_t padded_batch_size; + bool partition_kv; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { + return q_indptr[batch_idx + 1] - q_indptr[batch_idx]; + } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return kv_indptr[batch_idx + 1] - kv_indptr[batch_idx]; + } +}; + +struct PagedParams { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + DTypeQ* q; + paged_kv_t paged_kv; + IdType* q_indptr; + DTypeO* o; + float* lse; + uint_fastdiv group_size; + + uint8_t* maybe_custom_mask; +int32_t* maybe_mask_indptr; +float* maybe_alibi_slopes; +uint32_t* maybe_prefix_len_ptr; +uint16_t* maybe_token_pos_in_items_ptr; +uint16_t* maybe_max_item_len_ptr; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; +int64_t token_pos_in_items_len; + + uint32_t num_qo_heads; + IdType q_stride_n; + IdType q_stride_h; + int32_t window_left; + + IdType* request_indices; + IdType* qo_tile_indices; + IdType* kv_tile_indices; + IdType* merge_indptr; + IdType* o_indptr; + bool* block_valid_mask; + IdType* kv_chunk_size_ptr; + uint32_t max_total_num_rows; + uint32_t* total_num_rows; + uint32_t padded_batch_size; + bool partition_kv; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { + return q_indptr[batch_idx + 1] - q_indptr[batch_idx]; + } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return paged_kv.get_length(batch_idx); + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_jit_pybind.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..5421ab1cf910fdfbc867c182c8e22f23e44749ae --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_jit_pybind.cu @@ -0,0 +1,48 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "batch_prefill_config.inc" +#include "pytorch_extension_utils.h" + +at::Tensor BatchPrefillWithKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, + int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, + int64_t head_dim_vo, bool causal); + +void BatchPrefillWithRaggedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor k, at::Tensor v, + at::Tensor qo_indptr, at::Tensor kv_indptr, at::Tensor o, + std::optional maybe_lse, int64_t mask_mode_code, + int64_t layout, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +void BatchPrefillWithPagedKVCacheRun( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, + at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Batch-request prefill attention with KV-Cache plan + m.def("plan", BatchPrefillWithKVCachePlan); + // Batch-request prefill attention with KV-Cache operator + m.def("ragged_run", BatchPrefillWithRaggedKVCacheRun); + // Batch-request prefill attention with KV-Cache operator + m.def("paged_run", BatchPrefillWithPagedKVCacheRun); +} diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_paged_kernel_mask_0.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_paged_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..9ed3e16f02b81516d03b8e2b790df2742eb1430d --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_paged_kernel_mask_0.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kNone == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 256, 256, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 256, 256, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 256, 256, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_paged_kernel_mask_1.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_paged_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..37d4bfe8cdbda0809bc8b11730692da3632a1a44 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_paged_kernel_mask_1.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCausal == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 256, 256, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 256, 256, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 256, 256, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_paged_kernel_mask_2.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_paged_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..6d9b56ea582f9197cf55df1a48e80b7efce8273e --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_paged_kernel_mask_2.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCustom == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 256, 256, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 256, 256, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 256, 256, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_paged_kernel_mask_3.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_paged_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..4e67f0a63d0ba0aa6c2ef1be3f9f7e8e5bbac259 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_paged_kernel_mask_3.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kMultiItemScoring == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 256, 256, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 256, 256, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 256, 256, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_ragged_kernel_mask_0.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_ragged_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..4e4965782d385f94d63c616c3be1fdeef33474cb --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_ragged_kernel_mask_0.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kNone == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 256, 256, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 256, 256, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 256, 256, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_ragged_kernel_mask_1.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_ragged_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..eff49053b01f63913678acf1f2b86903f358de93 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_ragged_kernel_mask_1.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCausal == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 256, 256, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 256, 256, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 256, 256, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_ragged_kernel_mask_2.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_ragged_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..0c827cffe82db344b2afee6e33487871112ac672 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_ragged_kernel_mask_2.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCustom == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 256, 256, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 256, 256, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 256, 256, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_ragged_kernel_mask_3.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_ragged_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..5d7767dc6d8145fe1622713c8b1acbb1a61e4c8c --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_ragged_kernel_mask_3.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kMultiItemScoring == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 256, 256, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 256, 256, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 256, 256, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill.cu new file mode 100644 index 0000000000000000000000000000000000000000..a51fc7f56a0af7482bd5d2a9d3f672854df69dfd --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill.cu @@ -0,0 +1,330 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include +#include + +#include "batch_prefill_config.inc" +#include "pytorch_conversion_utils.h" +#include "pytorch_extension_utils.h" + +namespace flashinfer { + +template +cudaError_t BatchPrefillWithPagedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v, + float* tmp_s, bool enable_pdl, + cudaStream_t stream); + +template +cudaError_t BatchPrefillWithRaggedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v, + float* tmp_s, bool enable_pdl, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +at::Tensor BatchPrefillWithKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, + int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, + int64_t head_dim_vo, bool causal) { + size_t float_workspace_size_in_bytes = + float_workspace_buffer.size(0) * float_workspace_buffer.element_size(); + size_t int_workspace_size_in_bytes = + int_workspace_buffer.size(0) * int_workspace_buffer.element_size(); + + PrefillPlanInfo plan_info; + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + cudaError_t status = PrefillPlan( + float_workspace_buffer.data_ptr(), float_workspace_size_in_bytes, + int_workspace_buffer.data_ptr(), page_locked_int_workspace_buffer.data_ptr(), + int_workspace_size_in_bytes, plan_info, qo_indptr.data_ptr(), + kv_indptr.data_ptr(), total_num_rows, batch_size, num_qo_heads, num_kv_heads, + head_dim_qk, head_dim_vo, page_size, enable_cuda_graph, /*sizeof_dtype_o=*/2, stream); + + TORCH_CHECK(status == cudaSuccess, + "Failed to plan prefill with error: ", cudaGetErrorString(status)); + + return vec_to_tensor(plan_info.ToVector()); +} + +void BatchPrefillWithRaggedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor k, at::Tensor v, + at::Tensor qo_indptr, at::Tensor kv_indptr, at::Tensor o, + std::optional maybe_lse, int64_t mask_mode_code, + int64_t layout, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + PrefillPlanInfo plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + QKVLayout kv_layout = static_cast(layout); + + int64_t num_qo_heads = q.size(1); + int64_t head_dim_qk = q.size(2); + int64_t num_kv_heads = (kv_layout == QKVLayout::kNHD) ? k.size(1) : k.size(0); + uint32_t q_stride_n = q.stride(0), q_stride_h = q.stride(1), k_stride_n, k_stride_h, v_stride_n, + v_stride_h; + if (kv_layout == QKVLayout::kNHD) { + k_stride_n = k.stride(0); + k_stride_h = k.stride(1); + v_stride_n = v.stride(0); + v_stride_h = v.stride(1); + } else { + k_stride_h = k.stride(0); + k_stride_n = k.stride(1); + v_stride_h = v.stride(0); + v_stride_n = v.stride(1); + } + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == q.size(0), lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == q.size(1), lse.size(1), q.size(1)); + } + + void* float_buffer_ptr = float_workspace_buffer.data_ptr(); + void* int_buffer_ptr = int_workspace_buffer.data_ptr(); + + const MaskMode mask_mode = static_cast(mask_mode_code); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, + RaggedParams, PagedParams, [&] { + RaggedParams params; + + params.q = static_cast(q.data_ptr()); + params.k = static_cast(k.data_ptr()); + params.v = static_cast(v.data_ptr()); + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.q_indptr = static_cast(qo_indptr.data_ptr()); + params.kv_indptr = static_cast(kv_indptr.data_ptr()); + params.num_qo_heads = num_qo_heads; + params.num_kv_heads = num_kv_heads; + params.group_size = uint_fastdiv(num_qo_heads / num_kv_heads); + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.k_stride_n = k_stride_n; + params.k_stride_h = k_stride_h; + params.v_stride_n = v_stride_n; + params.v_stride_h = v_stride_h; + params.window_left = window_left; + + params.request_indices = nullptr; + params.qo_tile_indices = nullptr; + params.kv_tile_indices = nullptr; + params.merge_indptr = nullptr; + params.o_indptr = nullptr; + params.kv_chunk_size_ptr = nullptr; + params.block_valid_mask = nullptr; + params.total_num_rows = nullptr; + params.max_total_num_rows = 0; + params.padded_batch_size = 0; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + DTypeO* tmp_v = nullptr; + float* tmp_s = nullptr; + + params.request_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.request_indices_offset); + params.qo_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_tile_indices_offset); + params.kv_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_tile_indices_offset); + params.o_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.o_indptr_offset); + params.kv_chunk_size_ptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_chunk_size_ptr_offset); + if (plan_info.split_kv) { + params.merge_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.merge_indptr_offset); + tmp_v = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.v_offset); + tmp_s = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.s_offset); + if (plan_info.enable_cuda_graph) { + params.block_valid_mask = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.block_valid_mask_offset); + } + } + params.padded_batch_size = plan_info.padded_batch_size; + params.max_total_num_rows = plan_info.total_num_rows; + if (plan_info.enable_cuda_graph) { + params.total_num_rows = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.total_num_rows_offset); + } + + cudaError_t status = cudaSuccess; + + DISPATCH_CTA_TILE_Q(plan_info.cta_tile_q, CTA_TILE_Q, { + status = flashinfer::BatchPrefillWithRaggedKVCacheDispatched< + CTA_TILE_Q, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + /*use_fp16_qk_reduction=*/USE_FP16_QK_REDUCTION, MASK_MODE, AttentionVariant, + RaggedParams>(params, tmp_v, tmp_s, enable_pdl, stream); + }); + + TORCH_CHECK(status == cudaSuccess, "BatchPrefillWithRaggedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); +} + +void BatchPrefillWithPagedKVCacheRun( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, + at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + PrefillPlanInfo plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + QKVLayout kv_layout = static_cast(layout); + auto device = q.device(); + int64_t batch_size = paged_kv_indptr.size(0) - 1; + int64_t num_qo_heads = q.size(1); + int64_t num_kv_heads, page_size; + uint32_t head_dim_qk = q.size(2); + if (kv_layout == QKVLayout::kHND) { + num_kv_heads = paged_k_cache.size(1); + page_size = paged_k_cache.size(2); + } else { + page_size = paged_k_cache.size(1); + num_kv_heads = paged_k_cache.size(2); + } + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == q.size(0), lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == q.size(1), lse.size(1), q.size(1)); + } + + void* float_buffer_ptr = static_cast(float_workspace_buffer.data_ptr()); + void* int_buffer_ptr = static_cast(int_workspace_buffer.data_ptr()); + + const MaskMode mask_mode = static_cast(mask_mode_code); + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = paged_k_cache.scalar_type(); + + // get q_stride_n and q_stride_h + const auto q_stride_n = q.stride(0); + const auto q_stride_h = q.stride(1); + + // get kv_cache_strides + const int64_t* kv_cache_strides = nullptr; + auto k_strides = paged_k_cache.strides(); + auto v_strides = paged_v_cache.strides(); + TORCH_CHECK(k_strides == v_strides, "k/v strides must be identical"); + kv_cache_strides = k_strides.data(); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, + RaggedParams, PagedParams, [&] { + PagedParams params; + + params.q = static_cast(q.data_ptr()); + paged_kv_t paged_kv( + num_kv_heads, page_size, HEAD_DIM_VO, batch_size, kv_layout, + static_cast(paged_k_cache.data_ptr()), + static_cast(paged_v_cache.data_ptr()), kv_cache_strides, + static_cast(paged_kv_indices.data_ptr()), + static_cast(paged_kv_indptr.data_ptr()), + static_cast(paged_kv_last_page_len.data_ptr())); + params.paged_kv = paged_kv; + params.q_indptr = static_cast(qo_indptr.data_ptr()); + params.o = static_cast(o.data_ptr()); + + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.num_qo_heads = num_qo_heads; + params.group_size = uint_fastdiv(num_qo_heads / paged_kv.num_heads); + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.window_left = window_left; + + params.request_indices = nullptr; + params.qo_tile_indices = nullptr; + params.kv_tile_indices = nullptr; + params.merge_indptr = nullptr; + params.o_indptr = nullptr; + params.kv_chunk_size_ptr = nullptr; + params.block_valid_mask = nullptr; + params.total_num_rows = nullptr; + params.max_total_num_rows = 0; + params.padded_batch_size = 0; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + DTypeO* tmp_v = nullptr; + float* tmp_s = nullptr; + + params.request_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.request_indices_offset); + params.qo_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_tile_indices_offset); + params.kv_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_tile_indices_offset); + params.o_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.o_indptr_offset); + params.kv_chunk_size_ptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_chunk_size_ptr_offset); + if (plan_info.split_kv) { + params.merge_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.merge_indptr_offset); + tmp_v = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.v_offset); + tmp_s = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.s_offset); + if (plan_info.enable_cuda_graph) { + params.block_valid_mask = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.block_valid_mask_offset); + } + } + params.padded_batch_size = plan_info.padded_batch_size; + params.max_total_num_rows = plan_info.total_num_rows; + if (plan_info.enable_cuda_graph) { + params.total_num_rows = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.total_num_rows_offset); + } + + cudaError_t status = cudaSuccess; + + DISPATCH_CTA_TILE_Q(plan_info.cta_tile_q, CTA_TILE_Q, { + status = flashinfer::BatchPrefillWithPagedKVCacheDispatched< + CTA_TILE_Q, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + /*use_fp16_qk_reduction=*/USE_FP16_QK_REDUCTION, MASK_MODE, AttentionVariant, + PagedParams>(params, tmp_v, tmp_s, enable_pdl, stream); + }); + + TORCH_CHECK(status == cudaSuccess, "BatchPrefillWithPagedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); +} diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_config.inc b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..63127946eff38cc94d125a7d47e3a42b6831c1ac --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_config.inc @@ -0,0 +1,153 @@ +#pragma once +#include +#include +#include +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_custom_mask, std::optional maybe_mask_indptr, std::optional maybe_alibi_slopes, std::optional maybe_prefix_len_ptr, std::optional maybe_token_pos_in_items_ptr, std::optional maybe_max_item_len_ptr, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta, int64_t token_pos_in_items_len +#define ADDITIONAL_PARAMS_SETTER params.maybe_custom_mask = maybe_custom_mask ? static_cast(maybe_custom_mask->data_ptr()): nullptr; \ +params.maybe_mask_indptr = maybe_mask_indptr ? static_cast(maybe_mask_indptr->data_ptr()): nullptr; \ +params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.maybe_prefix_len_ptr = maybe_prefix_len_ptr ? static_cast(maybe_prefix_len_ptr->data_ptr()): nullptr; \ +params.maybe_token_pos_in_items_ptr = maybe_token_pos_in_items_ptr ? static_cast(maybe_token_pos_in_items_ptr->data_ptr()): nullptr; \ +params.maybe_max_item_len_ptr = maybe_max_item_len_ptr ? static_cast(maybe_max_item_len_ptr->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; \ +params.token_pos_in_items_len = token_pos_in_items_len; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, RaggedParams, PagedParams, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { \ + constexpr auto use_custom_mask = MASK_MODE == MaskMode::kCustom; \ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ + }) + +using namespace flashinfer; + +using DTypeQ = half; +using DTypeKV = __nv_fp8_e4m3; +using DTypeO = half; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 64; +constexpr int HEAD_DIM_VO = 64; +constexpr bool USE_FP16_QK_REDUCTION = false; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = false; + + +struct RaggedParams { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + DTypeQ* q; + DTypeKV* k; + DTypeKV* v; + IdType* q_indptr; + IdType* kv_indptr; + DTypeO* o; + float* lse; + uint_fastdiv group_size; + + uint8_t* maybe_custom_mask; +int32_t* maybe_mask_indptr; +float* maybe_alibi_slopes; +uint32_t* maybe_prefix_len_ptr; +uint16_t* maybe_token_pos_in_items_ptr; +uint16_t* maybe_max_item_len_ptr; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; +int64_t token_pos_in_items_len; + + uint32_t num_qo_heads; + uint32_t num_kv_heads; + uint32_t q_stride_n; + uint32_t q_stride_h; + uint32_t k_stride_n; + uint32_t k_stride_h; + uint32_t v_stride_n; + uint32_t v_stride_h; + int32_t window_left; + + IdType* request_indices; + IdType* qo_tile_indices; + IdType* kv_tile_indices; + IdType* merge_indptr; + IdType* o_indptr; + IdType* kv_chunk_size_ptr; + bool* block_valid_mask; + uint32_t max_total_num_rows; + uint32_t* total_num_rows; + uint32_t padded_batch_size; + bool partition_kv; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { + return q_indptr[batch_idx + 1] - q_indptr[batch_idx]; + } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return kv_indptr[batch_idx + 1] - kv_indptr[batch_idx]; + } +}; + +struct PagedParams { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + DTypeQ* q; + paged_kv_t paged_kv; + IdType* q_indptr; + DTypeO* o; + float* lse; + uint_fastdiv group_size; + + uint8_t* maybe_custom_mask; +int32_t* maybe_mask_indptr; +float* maybe_alibi_slopes; +uint32_t* maybe_prefix_len_ptr; +uint16_t* maybe_token_pos_in_items_ptr; +uint16_t* maybe_max_item_len_ptr; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; +int64_t token_pos_in_items_len; + + uint32_t num_qo_heads; + IdType q_stride_n; + IdType q_stride_h; + int32_t window_left; + + IdType* request_indices; + IdType* qo_tile_indices; + IdType* kv_tile_indices; + IdType* merge_indptr; + IdType* o_indptr; + bool* block_valid_mask; + IdType* kv_chunk_size_ptr; + uint32_t max_total_num_rows; + uint32_t* total_num_rows; + uint32_t padded_batch_size; + bool partition_kv; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { + return q_indptr[batch_idx + 1] - q_indptr[batch_idx]; + } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return paged_kv.get_length(batch_idx); + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_jit_pybind.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..5421ab1cf910fdfbc867c182c8e22f23e44749ae --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_jit_pybind.cu @@ -0,0 +1,48 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "batch_prefill_config.inc" +#include "pytorch_extension_utils.h" + +at::Tensor BatchPrefillWithKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, + int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, + int64_t head_dim_vo, bool causal); + +void BatchPrefillWithRaggedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor k, at::Tensor v, + at::Tensor qo_indptr, at::Tensor kv_indptr, at::Tensor o, + std::optional maybe_lse, int64_t mask_mode_code, + int64_t layout, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +void BatchPrefillWithPagedKVCacheRun( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, + at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Batch-request prefill attention with KV-Cache plan + m.def("plan", BatchPrefillWithKVCachePlan); + // Batch-request prefill attention with KV-Cache operator + m.def("ragged_run", BatchPrefillWithRaggedKVCacheRun); + // Batch-request prefill attention with KV-Cache operator + m.def("paged_run", BatchPrefillWithPagedKVCacheRun); +} diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_0.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..0d1186c56489ce15e29d5f1f07d672cd68998352 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_0.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kNone == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_1.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..b155d2ddc8021119a99d691c9a5f60762a7389ea --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_1.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCausal == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_2.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..344bc67ac9289dcf2830826f8e82b2bfac986d3d --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_2.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCustom == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_3.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..ae7ea93585a9b50effde7e743ac9fe2322864644 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_3.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kMultiItemScoring == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_0.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..63093bf106f131fe01fc68450a946b34192bdbbe --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_0.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kNone == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_1.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..dc4ad3d120d3eefac6d2c821b1bf0462549b455f --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_1.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCausal == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_2.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..7c289d9f839d07256ffb3a6cdff27774130e28c9 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_2.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCustom == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_3.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..8df8e96108f270774806db09caefab6836b5446b --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_3.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kMultiItemScoring == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill.cu new file mode 100644 index 0000000000000000000000000000000000000000..a51fc7f56a0af7482bd5d2a9d3f672854df69dfd --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill.cu @@ -0,0 +1,330 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include +#include + +#include "batch_prefill_config.inc" +#include "pytorch_conversion_utils.h" +#include "pytorch_extension_utils.h" + +namespace flashinfer { + +template +cudaError_t BatchPrefillWithPagedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v, + float* tmp_s, bool enable_pdl, + cudaStream_t stream); + +template +cudaError_t BatchPrefillWithRaggedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v, + float* tmp_s, bool enable_pdl, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +at::Tensor BatchPrefillWithKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, + int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, + int64_t head_dim_vo, bool causal) { + size_t float_workspace_size_in_bytes = + float_workspace_buffer.size(0) * float_workspace_buffer.element_size(); + size_t int_workspace_size_in_bytes = + int_workspace_buffer.size(0) * int_workspace_buffer.element_size(); + + PrefillPlanInfo plan_info; + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + cudaError_t status = PrefillPlan( + float_workspace_buffer.data_ptr(), float_workspace_size_in_bytes, + int_workspace_buffer.data_ptr(), page_locked_int_workspace_buffer.data_ptr(), + int_workspace_size_in_bytes, plan_info, qo_indptr.data_ptr(), + kv_indptr.data_ptr(), total_num_rows, batch_size, num_qo_heads, num_kv_heads, + head_dim_qk, head_dim_vo, page_size, enable_cuda_graph, /*sizeof_dtype_o=*/2, stream); + + TORCH_CHECK(status == cudaSuccess, + "Failed to plan prefill with error: ", cudaGetErrorString(status)); + + return vec_to_tensor(plan_info.ToVector()); +} + +void BatchPrefillWithRaggedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor k, at::Tensor v, + at::Tensor qo_indptr, at::Tensor kv_indptr, at::Tensor o, + std::optional maybe_lse, int64_t mask_mode_code, + int64_t layout, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + PrefillPlanInfo plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + QKVLayout kv_layout = static_cast(layout); + + int64_t num_qo_heads = q.size(1); + int64_t head_dim_qk = q.size(2); + int64_t num_kv_heads = (kv_layout == QKVLayout::kNHD) ? k.size(1) : k.size(0); + uint32_t q_stride_n = q.stride(0), q_stride_h = q.stride(1), k_stride_n, k_stride_h, v_stride_n, + v_stride_h; + if (kv_layout == QKVLayout::kNHD) { + k_stride_n = k.stride(0); + k_stride_h = k.stride(1); + v_stride_n = v.stride(0); + v_stride_h = v.stride(1); + } else { + k_stride_h = k.stride(0); + k_stride_n = k.stride(1); + v_stride_h = v.stride(0); + v_stride_n = v.stride(1); + } + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == q.size(0), lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == q.size(1), lse.size(1), q.size(1)); + } + + void* float_buffer_ptr = float_workspace_buffer.data_ptr(); + void* int_buffer_ptr = int_workspace_buffer.data_ptr(); + + const MaskMode mask_mode = static_cast(mask_mode_code); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, + RaggedParams, PagedParams, [&] { + RaggedParams params; + + params.q = static_cast(q.data_ptr()); + params.k = static_cast(k.data_ptr()); + params.v = static_cast(v.data_ptr()); + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.q_indptr = static_cast(qo_indptr.data_ptr()); + params.kv_indptr = static_cast(kv_indptr.data_ptr()); + params.num_qo_heads = num_qo_heads; + params.num_kv_heads = num_kv_heads; + params.group_size = uint_fastdiv(num_qo_heads / num_kv_heads); + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.k_stride_n = k_stride_n; + params.k_stride_h = k_stride_h; + params.v_stride_n = v_stride_n; + params.v_stride_h = v_stride_h; + params.window_left = window_left; + + params.request_indices = nullptr; + params.qo_tile_indices = nullptr; + params.kv_tile_indices = nullptr; + params.merge_indptr = nullptr; + params.o_indptr = nullptr; + params.kv_chunk_size_ptr = nullptr; + params.block_valid_mask = nullptr; + params.total_num_rows = nullptr; + params.max_total_num_rows = 0; + params.padded_batch_size = 0; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + DTypeO* tmp_v = nullptr; + float* tmp_s = nullptr; + + params.request_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.request_indices_offset); + params.qo_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_tile_indices_offset); + params.kv_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_tile_indices_offset); + params.o_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.o_indptr_offset); + params.kv_chunk_size_ptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_chunk_size_ptr_offset); + if (plan_info.split_kv) { + params.merge_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.merge_indptr_offset); + tmp_v = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.v_offset); + tmp_s = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.s_offset); + if (plan_info.enable_cuda_graph) { + params.block_valid_mask = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.block_valid_mask_offset); + } + } + params.padded_batch_size = plan_info.padded_batch_size; + params.max_total_num_rows = plan_info.total_num_rows; + if (plan_info.enable_cuda_graph) { + params.total_num_rows = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.total_num_rows_offset); + } + + cudaError_t status = cudaSuccess; + + DISPATCH_CTA_TILE_Q(plan_info.cta_tile_q, CTA_TILE_Q, { + status = flashinfer::BatchPrefillWithRaggedKVCacheDispatched< + CTA_TILE_Q, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + /*use_fp16_qk_reduction=*/USE_FP16_QK_REDUCTION, MASK_MODE, AttentionVariant, + RaggedParams>(params, tmp_v, tmp_s, enable_pdl, stream); + }); + + TORCH_CHECK(status == cudaSuccess, "BatchPrefillWithRaggedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); +} + +void BatchPrefillWithPagedKVCacheRun( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, + at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + PrefillPlanInfo plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + QKVLayout kv_layout = static_cast(layout); + auto device = q.device(); + int64_t batch_size = paged_kv_indptr.size(0) - 1; + int64_t num_qo_heads = q.size(1); + int64_t num_kv_heads, page_size; + uint32_t head_dim_qk = q.size(2); + if (kv_layout == QKVLayout::kHND) { + num_kv_heads = paged_k_cache.size(1); + page_size = paged_k_cache.size(2); + } else { + page_size = paged_k_cache.size(1); + num_kv_heads = paged_k_cache.size(2); + } + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == q.size(0), lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == q.size(1), lse.size(1), q.size(1)); + } + + void* float_buffer_ptr = static_cast(float_workspace_buffer.data_ptr()); + void* int_buffer_ptr = static_cast(int_workspace_buffer.data_ptr()); + + const MaskMode mask_mode = static_cast(mask_mode_code); + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = paged_k_cache.scalar_type(); + + // get q_stride_n and q_stride_h + const auto q_stride_n = q.stride(0); + const auto q_stride_h = q.stride(1); + + // get kv_cache_strides + const int64_t* kv_cache_strides = nullptr; + auto k_strides = paged_k_cache.strides(); + auto v_strides = paged_v_cache.strides(); + TORCH_CHECK(k_strides == v_strides, "k/v strides must be identical"); + kv_cache_strides = k_strides.data(); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, + RaggedParams, PagedParams, [&] { + PagedParams params; + + params.q = static_cast(q.data_ptr()); + paged_kv_t paged_kv( + num_kv_heads, page_size, HEAD_DIM_VO, batch_size, kv_layout, + static_cast(paged_k_cache.data_ptr()), + static_cast(paged_v_cache.data_ptr()), kv_cache_strides, + static_cast(paged_kv_indices.data_ptr()), + static_cast(paged_kv_indptr.data_ptr()), + static_cast(paged_kv_last_page_len.data_ptr())); + params.paged_kv = paged_kv; + params.q_indptr = static_cast(qo_indptr.data_ptr()); + params.o = static_cast(o.data_ptr()); + + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.num_qo_heads = num_qo_heads; + params.group_size = uint_fastdiv(num_qo_heads / paged_kv.num_heads); + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.window_left = window_left; + + params.request_indices = nullptr; + params.qo_tile_indices = nullptr; + params.kv_tile_indices = nullptr; + params.merge_indptr = nullptr; + params.o_indptr = nullptr; + params.kv_chunk_size_ptr = nullptr; + params.block_valid_mask = nullptr; + params.total_num_rows = nullptr; + params.max_total_num_rows = 0; + params.padded_batch_size = 0; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + DTypeO* tmp_v = nullptr; + float* tmp_s = nullptr; + + params.request_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.request_indices_offset); + params.qo_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_tile_indices_offset); + params.kv_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_tile_indices_offset); + params.o_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.o_indptr_offset); + params.kv_chunk_size_ptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_chunk_size_ptr_offset); + if (plan_info.split_kv) { + params.merge_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.merge_indptr_offset); + tmp_v = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.v_offset); + tmp_s = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.s_offset); + if (plan_info.enable_cuda_graph) { + params.block_valid_mask = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.block_valid_mask_offset); + } + } + params.padded_batch_size = plan_info.padded_batch_size; + params.max_total_num_rows = plan_info.total_num_rows; + if (plan_info.enable_cuda_graph) { + params.total_num_rows = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.total_num_rows_offset); + } + + cudaError_t status = cudaSuccess; + + DISPATCH_CTA_TILE_Q(plan_info.cta_tile_q, CTA_TILE_Q, { + status = flashinfer::BatchPrefillWithPagedKVCacheDispatched< + CTA_TILE_Q, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + /*use_fp16_qk_reduction=*/USE_FP16_QK_REDUCTION, MASK_MODE, AttentionVariant, + PagedParams>(params, tmp_v, tmp_s, enable_pdl, stream); + }); + + TORCH_CHECK(status == cudaSuccess, "BatchPrefillWithPagedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); +} diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_config.inc b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..ffc99697420f561c5d2abc4c997b50302ea2b5e9 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_config.inc @@ -0,0 +1,153 @@ +#pragma once +#include +#include +#include +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_custom_mask, std::optional maybe_mask_indptr, std::optional maybe_alibi_slopes, std::optional maybe_prefix_len_ptr, std::optional maybe_token_pos_in_items_ptr, std::optional maybe_max_item_len_ptr, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta, int64_t token_pos_in_items_len +#define ADDITIONAL_PARAMS_SETTER params.maybe_custom_mask = maybe_custom_mask ? static_cast(maybe_custom_mask->data_ptr()): nullptr; \ +params.maybe_mask_indptr = maybe_mask_indptr ? static_cast(maybe_mask_indptr->data_ptr()): nullptr; \ +params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.maybe_prefix_len_ptr = maybe_prefix_len_ptr ? static_cast(maybe_prefix_len_ptr->data_ptr()): nullptr; \ +params.maybe_token_pos_in_items_ptr = maybe_token_pos_in_items_ptr ? static_cast(maybe_token_pos_in_items_ptr->data_ptr()): nullptr; \ +params.maybe_max_item_len_ptr = maybe_max_item_len_ptr ? static_cast(maybe_max_item_len_ptr->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; \ +params.token_pos_in_items_len = token_pos_in_items_len; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, RaggedParams, PagedParams, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { \ + constexpr auto use_custom_mask = MASK_MODE == MaskMode::kCustom; \ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ + }) + +using namespace flashinfer; + +using DTypeQ = half; +using DTypeKV = __nv_fp8_e4m3; +using DTypeO = half; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 64; +constexpr int HEAD_DIM_VO = 64; +constexpr bool USE_FP16_QK_REDUCTION = false; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = true; + + +struct RaggedParams { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + DTypeQ* q; + DTypeKV* k; + DTypeKV* v; + IdType* q_indptr; + IdType* kv_indptr; + DTypeO* o; + float* lse; + uint_fastdiv group_size; + + uint8_t* maybe_custom_mask; +int32_t* maybe_mask_indptr; +float* maybe_alibi_slopes; +uint32_t* maybe_prefix_len_ptr; +uint16_t* maybe_token_pos_in_items_ptr; +uint16_t* maybe_max_item_len_ptr; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; +int64_t token_pos_in_items_len; + + uint32_t num_qo_heads; + uint32_t num_kv_heads; + uint32_t q_stride_n; + uint32_t q_stride_h; + uint32_t k_stride_n; + uint32_t k_stride_h; + uint32_t v_stride_n; + uint32_t v_stride_h; + int32_t window_left; + + IdType* request_indices; + IdType* qo_tile_indices; + IdType* kv_tile_indices; + IdType* merge_indptr; + IdType* o_indptr; + IdType* kv_chunk_size_ptr; + bool* block_valid_mask; + uint32_t max_total_num_rows; + uint32_t* total_num_rows; + uint32_t padded_batch_size; + bool partition_kv; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { + return q_indptr[batch_idx + 1] - q_indptr[batch_idx]; + } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return kv_indptr[batch_idx + 1] - kv_indptr[batch_idx]; + } +}; + +struct PagedParams { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + DTypeQ* q; + paged_kv_t paged_kv; + IdType* q_indptr; + DTypeO* o; + float* lse; + uint_fastdiv group_size; + + uint8_t* maybe_custom_mask; +int32_t* maybe_mask_indptr; +float* maybe_alibi_slopes; +uint32_t* maybe_prefix_len_ptr; +uint16_t* maybe_token_pos_in_items_ptr; +uint16_t* maybe_max_item_len_ptr; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; +int64_t token_pos_in_items_len; + + uint32_t num_qo_heads; + IdType q_stride_n; + IdType q_stride_h; + int32_t window_left; + + IdType* request_indices; + IdType* qo_tile_indices; + IdType* kv_tile_indices; + IdType* merge_indptr; + IdType* o_indptr; + bool* block_valid_mask; + IdType* kv_chunk_size_ptr; + uint32_t max_total_num_rows; + uint32_t* total_num_rows; + uint32_t padded_batch_size; + bool partition_kv; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { + return q_indptr[batch_idx + 1] - q_indptr[batch_idx]; + } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return paged_kv.get_length(batch_idx); + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_jit_pybind.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..5421ab1cf910fdfbc867c182c8e22f23e44749ae --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_jit_pybind.cu @@ -0,0 +1,48 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "batch_prefill_config.inc" +#include "pytorch_extension_utils.h" + +at::Tensor BatchPrefillWithKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, + int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, + int64_t head_dim_vo, bool causal); + +void BatchPrefillWithRaggedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor k, at::Tensor v, + at::Tensor qo_indptr, at::Tensor kv_indptr, at::Tensor o, + std::optional maybe_lse, int64_t mask_mode_code, + int64_t layout, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +void BatchPrefillWithPagedKVCacheRun( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, + at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Batch-request prefill attention with KV-Cache plan + m.def("plan", BatchPrefillWithKVCachePlan); + // Batch-request prefill attention with KV-Cache operator + m.def("ragged_run", BatchPrefillWithRaggedKVCacheRun); + // Batch-request prefill attention with KV-Cache operator + m.def("paged_run", BatchPrefillWithPagedKVCacheRun); +} diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_0.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..54b583ba709acb7db72c4e1854f934ef3fc54872 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_0.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kNone == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_1.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..ffceb828ae6798d6d4cd55042d9578f9924157b3 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_1.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCausal == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_2.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..1e52ae59b60c66d2eb32d4ec8c1eb829c9153df0 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_2.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCustom == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_3.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..170d1a6e67e5d2791d29cdc46d28e6755a237e44 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_3.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kMultiItemScoring == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_0.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..96536df05b7b979cf5f79585eb7c4c9ae5ba2bbc --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_0.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kNone == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_1.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..8c513ea2d6fdec8108b49c8e94da978831075ddb --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_1.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCausal == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_2.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..b2e25b8da455088c2e73d4002cbf78824c0dfc39 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_2.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCustom == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_3.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..6dd8e0b2a9973c9b006547e05016f1eba273a485 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_3.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kMultiItemScoring == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill.cu new file mode 100644 index 0000000000000000000000000000000000000000..a51fc7f56a0af7482bd5d2a9d3f672854df69dfd --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill.cu @@ -0,0 +1,330 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include +#include + +#include "batch_prefill_config.inc" +#include "pytorch_conversion_utils.h" +#include "pytorch_extension_utils.h" + +namespace flashinfer { + +template +cudaError_t BatchPrefillWithPagedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v, + float* tmp_s, bool enable_pdl, + cudaStream_t stream); + +template +cudaError_t BatchPrefillWithRaggedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v, + float* tmp_s, bool enable_pdl, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +at::Tensor BatchPrefillWithKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, + int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, + int64_t head_dim_vo, bool causal) { + size_t float_workspace_size_in_bytes = + float_workspace_buffer.size(0) * float_workspace_buffer.element_size(); + size_t int_workspace_size_in_bytes = + int_workspace_buffer.size(0) * int_workspace_buffer.element_size(); + + PrefillPlanInfo plan_info; + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + cudaError_t status = PrefillPlan( + float_workspace_buffer.data_ptr(), float_workspace_size_in_bytes, + int_workspace_buffer.data_ptr(), page_locked_int_workspace_buffer.data_ptr(), + int_workspace_size_in_bytes, plan_info, qo_indptr.data_ptr(), + kv_indptr.data_ptr(), total_num_rows, batch_size, num_qo_heads, num_kv_heads, + head_dim_qk, head_dim_vo, page_size, enable_cuda_graph, /*sizeof_dtype_o=*/2, stream); + + TORCH_CHECK(status == cudaSuccess, + "Failed to plan prefill with error: ", cudaGetErrorString(status)); + + return vec_to_tensor(plan_info.ToVector()); +} + +void BatchPrefillWithRaggedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor k, at::Tensor v, + at::Tensor qo_indptr, at::Tensor kv_indptr, at::Tensor o, + std::optional maybe_lse, int64_t mask_mode_code, + int64_t layout, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + PrefillPlanInfo plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + QKVLayout kv_layout = static_cast(layout); + + int64_t num_qo_heads = q.size(1); + int64_t head_dim_qk = q.size(2); + int64_t num_kv_heads = (kv_layout == QKVLayout::kNHD) ? k.size(1) : k.size(0); + uint32_t q_stride_n = q.stride(0), q_stride_h = q.stride(1), k_stride_n, k_stride_h, v_stride_n, + v_stride_h; + if (kv_layout == QKVLayout::kNHD) { + k_stride_n = k.stride(0); + k_stride_h = k.stride(1); + v_stride_n = v.stride(0); + v_stride_h = v.stride(1); + } else { + k_stride_h = k.stride(0); + k_stride_n = k.stride(1); + v_stride_h = v.stride(0); + v_stride_n = v.stride(1); + } + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == q.size(0), lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == q.size(1), lse.size(1), q.size(1)); + } + + void* float_buffer_ptr = float_workspace_buffer.data_ptr(); + void* int_buffer_ptr = int_workspace_buffer.data_ptr(); + + const MaskMode mask_mode = static_cast(mask_mode_code); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, + RaggedParams, PagedParams, [&] { + RaggedParams params; + + params.q = static_cast(q.data_ptr()); + params.k = static_cast(k.data_ptr()); + params.v = static_cast(v.data_ptr()); + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.q_indptr = static_cast(qo_indptr.data_ptr()); + params.kv_indptr = static_cast(kv_indptr.data_ptr()); + params.num_qo_heads = num_qo_heads; + params.num_kv_heads = num_kv_heads; + params.group_size = uint_fastdiv(num_qo_heads / num_kv_heads); + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.k_stride_n = k_stride_n; + params.k_stride_h = k_stride_h; + params.v_stride_n = v_stride_n; + params.v_stride_h = v_stride_h; + params.window_left = window_left; + + params.request_indices = nullptr; + params.qo_tile_indices = nullptr; + params.kv_tile_indices = nullptr; + params.merge_indptr = nullptr; + params.o_indptr = nullptr; + params.kv_chunk_size_ptr = nullptr; + params.block_valid_mask = nullptr; + params.total_num_rows = nullptr; + params.max_total_num_rows = 0; + params.padded_batch_size = 0; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + DTypeO* tmp_v = nullptr; + float* tmp_s = nullptr; + + params.request_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.request_indices_offset); + params.qo_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_tile_indices_offset); + params.kv_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_tile_indices_offset); + params.o_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.o_indptr_offset); + params.kv_chunk_size_ptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_chunk_size_ptr_offset); + if (plan_info.split_kv) { + params.merge_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.merge_indptr_offset); + tmp_v = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.v_offset); + tmp_s = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.s_offset); + if (plan_info.enable_cuda_graph) { + params.block_valid_mask = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.block_valid_mask_offset); + } + } + params.padded_batch_size = plan_info.padded_batch_size; + params.max_total_num_rows = plan_info.total_num_rows; + if (plan_info.enable_cuda_graph) { + params.total_num_rows = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.total_num_rows_offset); + } + + cudaError_t status = cudaSuccess; + + DISPATCH_CTA_TILE_Q(plan_info.cta_tile_q, CTA_TILE_Q, { + status = flashinfer::BatchPrefillWithRaggedKVCacheDispatched< + CTA_TILE_Q, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + /*use_fp16_qk_reduction=*/USE_FP16_QK_REDUCTION, MASK_MODE, AttentionVariant, + RaggedParams>(params, tmp_v, tmp_s, enable_pdl, stream); + }); + + TORCH_CHECK(status == cudaSuccess, "BatchPrefillWithRaggedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); +} + +void BatchPrefillWithPagedKVCacheRun( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, + at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + PrefillPlanInfo plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + QKVLayout kv_layout = static_cast(layout); + auto device = q.device(); + int64_t batch_size = paged_kv_indptr.size(0) - 1; + int64_t num_qo_heads = q.size(1); + int64_t num_kv_heads, page_size; + uint32_t head_dim_qk = q.size(2); + if (kv_layout == QKVLayout::kHND) { + num_kv_heads = paged_k_cache.size(1); + page_size = paged_k_cache.size(2); + } else { + page_size = paged_k_cache.size(1); + num_kv_heads = paged_k_cache.size(2); + } + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == q.size(0), lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == q.size(1), lse.size(1), q.size(1)); + } + + void* float_buffer_ptr = static_cast(float_workspace_buffer.data_ptr()); + void* int_buffer_ptr = static_cast(int_workspace_buffer.data_ptr()); + + const MaskMode mask_mode = static_cast(mask_mode_code); + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = paged_k_cache.scalar_type(); + + // get q_stride_n and q_stride_h + const auto q_stride_n = q.stride(0); + const auto q_stride_h = q.stride(1); + + // get kv_cache_strides + const int64_t* kv_cache_strides = nullptr; + auto k_strides = paged_k_cache.strides(); + auto v_strides = paged_v_cache.strides(); + TORCH_CHECK(k_strides == v_strides, "k/v strides must be identical"); + kv_cache_strides = k_strides.data(); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, + RaggedParams, PagedParams, [&] { + PagedParams params; + + params.q = static_cast(q.data_ptr()); + paged_kv_t paged_kv( + num_kv_heads, page_size, HEAD_DIM_VO, batch_size, kv_layout, + static_cast(paged_k_cache.data_ptr()), + static_cast(paged_v_cache.data_ptr()), kv_cache_strides, + static_cast(paged_kv_indices.data_ptr()), + static_cast(paged_kv_indptr.data_ptr()), + static_cast(paged_kv_last_page_len.data_ptr())); + params.paged_kv = paged_kv; + params.q_indptr = static_cast(qo_indptr.data_ptr()); + params.o = static_cast(o.data_ptr()); + + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.num_qo_heads = num_qo_heads; + params.group_size = uint_fastdiv(num_qo_heads / paged_kv.num_heads); + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.window_left = window_left; + + params.request_indices = nullptr; + params.qo_tile_indices = nullptr; + params.kv_tile_indices = nullptr; + params.merge_indptr = nullptr; + params.o_indptr = nullptr; + params.kv_chunk_size_ptr = nullptr; + params.block_valid_mask = nullptr; + params.total_num_rows = nullptr; + params.max_total_num_rows = 0; + params.padded_batch_size = 0; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + DTypeO* tmp_v = nullptr; + float* tmp_s = nullptr; + + params.request_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.request_indices_offset); + params.qo_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_tile_indices_offset); + params.kv_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_tile_indices_offset); + params.o_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.o_indptr_offset); + params.kv_chunk_size_ptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_chunk_size_ptr_offset); + if (plan_info.split_kv) { + params.merge_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.merge_indptr_offset); + tmp_v = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.v_offset); + tmp_s = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.s_offset); + if (plan_info.enable_cuda_graph) { + params.block_valid_mask = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.block_valid_mask_offset); + } + } + params.padded_batch_size = plan_info.padded_batch_size; + params.max_total_num_rows = plan_info.total_num_rows; + if (plan_info.enable_cuda_graph) { + params.total_num_rows = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.total_num_rows_offset); + } + + cudaError_t status = cudaSuccess; + + DISPATCH_CTA_TILE_Q(plan_info.cta_tile_q, CTA_TILE_Q, { + status = flashinfer::BatchPrefillWithPagedKVCacheDispatched< + CTA_TILE_Q, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + /*use_fp16_qk_reduction=*/USE_FP16_QK_REDUCTION, MASK_MODE, AttentionVariant, + PagedParams>(params, tmp_v, tmp_s, enable_pdl, stream); + }); + + TORCH_CHECK(status == cudaSuccess, "BatchPrefillWithPagedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); +} diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_config.inc b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..a628eb6c041713ec0d1307fdb65d023e27dfd96a --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_config.inc @@ -0,0 +1,153 @@ +#pragma once +#include +#include +#include +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_custom_mask, std::optional maybe_mask_indptr, std::optional maybe_alibi_slopes, std::optional maybe_prefix_len_ptr, std::optional maybe_token_pos_in_items_ptr, std::optional maybe_max_item_len_ptr, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta, int64_t token_pos_in_items_len +#define ADDITIONAL_PARAMS_SETTER params.maybe_custom_mask = maybe_custom_mask ? static_cast(maybe_custom_mask->data_ptr()): nullptr; \ +params.maybe_mask_indptr = maybe_mask_indptr ? static_cast(maybe_mask_indptr->data_ptr()): nullptr; \ +params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.maybe_prefix_len_ptr = maybe_prefix_len_ptr ? static_cast(maybe_prefix_len_ptr->data_ptr()): nullptr; \ +params.maybe_token_pos_in_items_ptr = maybe_token_pos_in_items_ptr ? static_cast(maybe_token_pos_in_items_ptr->data_ptr()): nullptr; \ +params.maybe_max_item_len_ptr = maybe_max_item_len_ptr ? static_cast(maybe_max_item_len_ptr->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; \ +params.token_pos_in_items_len = token_pos_in_items_len; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, RaggedParams, PagedParams, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { \ + constexpr auto use_custom_mask = MASK_MODE == MaskMode::kCustom; \ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ + }) + +using namespace flashinfer; + +using DTypeQ = half; +using DTypeKV = half; +using DTypeO = half; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 128; +constexpr int HEAD_DIM_VO = 128; +constexpr bool USE_FP16_QK_REDUCTION = false; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = false; + + +struct RaggedParams { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + DTypeQ* q; + DTypeKV* k; + DTypeKV* v; + IdType* q_indptr; + IdType* kv_indptr; + DTypeO* o; + float* lse; + uint_fastdiv group_size; + + uint8_t* maybe_custom_mask; +int32_t* maybe_mask_indptr; +float* maybe_alibi_slopes; +uint32_t* maybe_prefix_len_ptr; +uint16_t* maybe_token_pos_in_items_ptr; +uint16_t* maybe_max_item_len_ptr; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; +int64_t token_pos_in_items_len; + + uint32_t num_qo_heads; + uint32_t num_kv_heads; + uint32_t q_stride_n; + uint32_t q_stride_h; + uint32_t k_stride_n; + uint32_t k_stride_h; + uint32_t v_stride_n; + uint32_t v_stride_h; + int32_t window_left; + + IdType* request_indices; + IdType* qo_tile_indices; + IdType* kv_tile_indices; + IdType* merge_indptr; + IdType* o_indptr; + IdType* kv_chunk_size_ptr; + bool* block_valid_mask; + uint32_t max_total_num_rows; + uint32_t* total_num_rows; + uint32_t padded_batch_size; + bool partition_kv; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { + return q_indptr[batch_idx + 1] - q_indptr[batch_idx]; + } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return kv_indptr[batch_idx + 1] - kv_indptr[batch_idx]; + } +}; + +struct PagedParams { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + DTypeQ* q; + paged_kv_t paged_kv; + IdType* q_indptr; + DTypeO* o; + float* lse; + uint_fastdiv group_size; + + uint8_t* maybe_custom_mask; +int32_t* maybe_mask_indptr; +float* maybe_alibi_slopes; +uint32_t* maybe_prefix_len_ptr; +uint16_t* maybe_token_pos_in_items_ptr; +uint16_t* maybe_max_item_len_ptr; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; +int64_t token_pos_in_items_len; + + uint32_t num_qo_heads; + IdType q_stride_n; + IdType q_stride_h; + int32_t window_left; + + IdType* request_indices; + IdType* qo_tile_indices; + IdType* kv_tile_indices; + IdType* merge_indptr; + IdType* o_indptr; + bool* block_valid_mask; + IdType* kv_chunk_size_ptr; + uint32_t max_total_num_rows; + uint32_t* total_num_rows; + uint32_t padded_batch_size; + bool partition_kv; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { + return q_indptr[batch_idx + 1] - q_indptr[batch_idx]; + } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return paged_kv.get_length(batch_idx); + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_jit_pybind.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..5421ab1cf910fdfbc867c182c8e22f23e44749ae --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_jit_pybind.cu @@ -0,0 +1,48 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "batch_prefill_config.inc" +#include "pytorch_extension_utils.h" + +at::Tensor BatchPrefillWithKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, + int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, + int64_t head_dim_vo, bool causal); + +void BatchPrefillWithRaggedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor k, at::Tensor v, + at::Tensor qo_indptr, at::Tensor kv_indptr, at::Tensor o, + std::optional maybe_lse, int64_t mask_mode_code, + int64_t layout, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +void BatchPrefillWithPagedKVCacheRun( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, + at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Batch-request prefill attention with KV-Cache plan + m.def("plan", BatchPrefillWithKVCachePlan); + // Batch-request prefill attention with KV-Cache operator + m.def("ragged_run", BatchPrefillWithRaggedKVCacheRun); + // Batch-request prefill attention with KV-Cache operator + m.def("paged_run", BatchPrefillWithPagedKVCacheRun); +} diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_0.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..ff4fbd96e95fa670c019e4fb5cda473f41ba1628 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_0.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kNone == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 128, 128, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 128, 128, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 128, 128, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_1.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..ae257f99a6b4f7b44ff90ee4fb0395a49dedaf9c --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_1.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCausal == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 128, 128, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 128, 128, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 128, 128, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_2.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..9baac31703e1a9977eb8e85ef79bd23e4634ae5c --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_2.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCustom == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 128, 128, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 128, 128, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 128, 128, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_3.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..108509bf3bd71e251f388881bca3f7f7f26154e7 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_3.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kMultiItemScoring == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 128, 128, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 128, 128, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 128, 128, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_0.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..39450198451b3c42f30ed9e3eb829e0be30fe4a1 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_0.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kNone == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 128, 128, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 128, 128, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 128, 128, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_1.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..50ab7b2f475eedf92137e309509db05f309cc591 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_1.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCausal == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 128, 128, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 128, 128, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 128, 128, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_2.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..86408ffce4c329caabf1f856d054120815065520 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_2.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCustom == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 128, 128, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 128, 128, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 128, 128, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_3.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..b44ab24580d5632515cf18460b859352aa95754e --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_3.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kMultiItemScoring == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 128, 128, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 128, 128, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 128, 128, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_0.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..7a3dcb47ab89daa3289551969a12c2a2c53c3d19 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_0.cu @@ -0,0 +1,24 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched + <128, + 128, + MaskMode::kNone, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultAttention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched + <128, + 128, + MaskMode::kNone, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultAttention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_1.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..7374938760f2ca572a9bcee91465b15daf194e5d --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_1.cu @@ -0,0 +1,24 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched + <128, + 128, + MaskMode::kCausal, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultAttention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched + <128, + 128, + MaskMode::kCausal, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultAttention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_2.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..522184fb2d414cba8ce56e576d475fee0651f912 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_2.cu @@ -0,0 +1,24 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched + <128, + 128, + MaskMode::kCustom, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultAttention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched + <128, + 128, + MaskMode::kCustom, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultAttention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_3.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..4a657661a8d94cbdc5c4ee19bb75e13db7588a2b --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_3.cu @@ -0,0 +1,24 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched + <128, + 128, + MaskMode::kMultiItemScoring, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultAttention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched + <128, + 128, + MaskMode::kMultiItemScoring, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultAttention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_0.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..bf033b46b81da64fc48ade098472326bd5c2712f --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_0.cu @@ -0,0 +1,24 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched + <128, + 128, + MaskMode::kNone, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultAttention>(RaggedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched + <128, + 128, + MaskMode::kNone, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultAttention>(RaggedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_1.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..51e21ba99d82e9e4d3840ea6dc1ee260d7550580 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_1.cu @@ -0,0 +1,24 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched + <128, + 128, + MaskMode::kCausal, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultAttention>(RaggedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched + <128, + 128, + MaskMode::kCausal, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultAttention>(RaggedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_2.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..0c7e6867e3aed1752226dc283c266ed97682e290 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_2.cu @@ -0,0 +1,24 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched + <128, + 128, + MaskMode::kCustom, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultAttention>(RaggedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched + <128, + 128, + MaskMode::kCustom, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultAttention>(RaggedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_3.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..e15debd96911a4e77e702830b7c4d414410e3e98 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_3.cu @@ -0,0 +1,24 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched + <128, + 128, + MaskMode::kMultiItemScoring, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultAttention>(RaggedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched + <128, + 128, + MaskMode::kMultiItemScoring, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultAttention>(RaggedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_sm90.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_sm90.cu new file mode 100644 index 0000000000000000000000000000000000000000..159ee3f1eda3f47769ec14f8fac16f35fa9bda05 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_sm90.cu @@ -0,0 +1,260 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include +#include + +#include "batch_prefill_sm90_config.inc" +#include "pytorch_conversion_utils.h" +#include "pytorch_extension_utils.h" + +namespace flashinfer { + +template +cudaError_t BatchPrefillWithRaggedKVCacheDispatched(Params& params, bool enable_pdl, + cudaStream_t stream); + +template +cudaError_t BatchPrefillWithPagedKVCacheDispatched(Params& params, bool enable_pdl, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +at::Tensor BatchPrefillWithKVCacheSM90Plan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, + int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, + int64_t head_dim_vo, bool causal) { + size_t float_workspace_size_in_bytes = + float_workspace_buffer.size(0) * float_workspace_buffer.element_size(); + size_t int_workspace_size_in_bytes = + int_workspace_buffer.size(0) * int_workspace_buffer.element_size(); + + flashinfer::PrefillPlanSM90Info plan_info; + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + cudaError_t status = + PrefillSM90Plan(float_workspace_buffer.data_ptr(), float_workspace_size_in_bytes, + int_workspace_buffer.data_ptr(), page_locked_int_workspace_buffer.data_ptr(), + int_workspace_size_in_bytes, plan_info, qo_indptr.data_ptr(), + kv_indptr.data_ptr(), kv_len_arr.data_ptr(), total_num_rows, + batch_size, num_qo_heads, num_kv_heads, head_dim_qk, head_dim_vo, page_size, + causal, enable_cuda_graph, /*sizeof_dtype_o=*/2, stream); + + TORCH_CHECK(status == cudaSuccess, + "PrefillSM90Plan failed with error: ", cudaGetErrorString(status)); + + return vec_to_tensor(plan_info.ToVector()); +} + +void BatchPrefillWithRaggedKVCacheSM90Run( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + PrefillPlanSM90Info plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == q.size(0), lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == q.size(1), lse.size(1), q.size(1)); + } + + void* float_buffer_ptr = float_workspace_buffer.data_ptr(); + void* int_buffer_ptr = int_workspace_buffer.data_ptr(); + + int64_t head_dim_qk = q.size(2); + int64_t head_dim_vo = v.size(2); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + QKVLayout kv_layout = static_cast(layout); + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + const MaskMode mask_mode = static_cast(mask_mode_code); + bool use_swa = window_left != -1; + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, USE_SLIDING_WINDOW, + USE_LOGITS_SOFT_CAP, AttentionVariant, RaggedParams, PagedParams, [&] { + RaggedParams params; + + params.q_ptr = static_cast(q.data_ptr()); + params.k_ptr = static_cast(k.data_ptr()); + params.v_ptr = static_cast(v.data_ptr()); + params.o_ptr = static_cast(o.data_ptr()); + params.lse_ptr = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.q_stride_n = q.stride(0); + params.q_stride_h = q.stride(1); + params.o_stride_n = o.stride(0); + params.o_stride_h = o.stride(1); + if (kv_layout == QKVLayout::kNHD) { + params.k_stride_n = k.stride(0); + params.k_stride_h = k.stride(1); + params.v_stride_n = v.stride(0); + params.v_stride_h = v.stride(1); + } else { + params.k_stride_h = k.stride(0); + params.k_stride_n = k.stride(1); + params.v_stride_h = v.stride(0); + params.v_stride_n = v.stride(1); + } + params.nnz_qo = q.size(0); + params.nnz_kv = k.size(0); + params.num_qo_heads = q.size(1); + params.num_kv_heads = k.size(1); + params.group_size = params.num_qo_heads / params.num_kv_heads; + params.window_left = window_left; + params.causal = mask_mode_code == 1; + params.qo_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_tile_indices_offset); + params.qo_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_indptr_offset); + params.kv_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_indptr_offset); + params.qo_lens = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_len_offset); + params.kv_lens = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_len_offset); + params.head_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.head_indices_offset); + params.work_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.work_indptr_offset); + params.batch_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.batch_indices_offset); + + ADDITIONAL_PARAMS_SETTER + + bool same_schedule_for_all_heads = plan_info.same_schedule_for_all_heads; + DISPATCH_BOOL(same_schedule_for_all_heads, SAME_SCHEDULER_FOR_ALL_HEADS, [&] { + cudaError_t status = BatchPrefillWithRaggedKVCacheDispatched< + HEAD_DIM_QK, HEAD_DIM_VO, MASK_MODE, USE_SLIDING_WINDOW, SAME_SCHEDULER_FOR_ALL_HEADS, + AttentionVariant>(params, enable_pdl, stream); + TORCH_CHECK(status == cudaSuccess, + "BatchPrefillWithRaggedKVCacheSM90Run failed with error: ", + cudaGetErrorString(status)); + return true; + }); + }); +} + +void BatchPrefillWithPagedKVCacheSM90Run( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, + at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + PrefillPlanSM90Info plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == q.size(0), lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == q.size(1), lse.size(1), q.size(1)); + } + QKVLayout kv_layout = static_cast(layout); + int64_t num_kv_heads, page_size; + int64_t head_dim_qk = q.size(2); + int64_t head_dim_vo = paged_v_cache.size(3); + if (kv_layout == QKVLayout::kHND) { + num_kv_heads = paged_k_cache.size(1); + page_size = paged_k_cache.size(2); + } else { + page_size = paged_k_cache.size(1); + num_kv_heads = paged_k_cache.size(2); + } + + void* float_buffer_ptr = float_workspace_buffer.data_ptr(); + void* int_buffer_ptr = int_workspace_buffer.data_ptr(); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = paged_k_cache.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + const MaskMode mask_mode = static_cast(mask_mode_code); + bool use_swa = window_left != -1; + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, USE_SLIDING_WINDOW, + USE_LOGITS_SOFT_CAP, AttentionVariant, RaggedParams, PagedParams, [&] { + PagedParams params; + + params.q_ptr = static_cast(q.data_ptr()); + params.k_ptr = static_cast(paged_k_cache.data_ptr()); + params.v_ptr = static_cast(paged_v_cache.data_ptr()); + params.o_ptr = static_cast(o.data_ptr()); + params.lse_ptr = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.q_stride_n = q.stride(0); + params.q_stride_h = q.stride(1); + params.o_stride_n = o.stride(0); + params.o_stride_h = o.stride(1); + if (kv_layout == QKVLayout::kNHD) { + // (num_pages, page_size, num_heads, head_dim) + params.k_stride_n = paged_k_cache.stride(1); + params.k_stride_h = paged_k_cache.stride(2); + params.v_stride_n = paged_v_cache.stride(1); + params.v_stride_h = paged_v_cache.stride(2); + } else { + // (num_pages, num_heads, page_size, head_dim) + params.k_stride_h = paged_k_cache.stride(1); + params.k_stride_n = paged_k_cache.stride(2); + params.v_stride_h = paged_v_cache.stride(1); + params.v_stride_n = paged_v_cache.stride(2); + } + params.nnz_qo = q.size(0); + params.num_qo_heads = q.size(1); + params.num_kv_heads = num_kv_heads; + params.group_size = params.num_qo_heads / num_kv_heads; + params.page_size = page_size; + params.window_left = window_left; + params.causal = mask_mode_code == 1; + params.qo_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_tile_indices_offset); + params.qo_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_indptr_offset); + params.kv_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_indptr_offset); + params.qo_lens = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_len_offset); + params.kv_lens = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_len_offset); + params.head_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.head_indices_offset); + params.work_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.work_indptr_offset); + params.batch_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.batch_indices_offset); + params.kv_indices = static_cast(paged_kv_indices.data_ptr()); + + ADDITIONAL_PARAMS_SETTER + + bool same_schedule_for_all_heads = plan_info.same_schedule_for_all_heads; + DISPATCH_BOOL(same_schedule_for_all_heads, SAME_SCHEDULER_FOR_ALL_HEADS, [&] { + cudaError_t status = BatchPrefillWithPagedKVCacheDispatched< + HEAD_DIM_QK, HEAD_DIM_VO, MASK_MODE, USE_SLIDING_WINDOW, SAME_SCHEDULER_FOR_ALL_HEADS, + AttentionVariant>(params, enable_pdl, stream); + TORCH_CHECK(status == cudaSuccess, + "BatchPrefillWithPagedKVCacheSM90Run failed with error: ", + cudaGetErrorString(status)); + return true; + }); + }); +} diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_sm90_config.inc b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_sm90_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..6b18b676ac5f99b978a44cfa6ff1ca30ffa635a2 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_sm90_config.inc @@ -0,0 +1,135 @@ +#pragma once +#include +#include +#include +#include +#include +#include "pytorch_extension_utils.h" + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_prefix_len_ptr, std::optional maybe_token_pos_in_items_ptr, std::optional maybe_max_item_len_ptr, double logits_soft_cap, double sm_scale, int64_t token_pos_in_items_len +#define ADDITIONAL_PARAMS_SETTER params.additional_params.maybe_prefix_len_ptr = maybe_prefix_len_ptr ? static_cast(maybe_prefix_len_ptr->data_ptr()): nullptr; \ +params.additional_params.maybe_token_pos_in_items_ptr = maybe_token_pos_in_items_ptr ? static_cast(maybe_token_pos_in_items_ptr->data_ptr()): nullptr; \ +params.additional_params.maybe_max_item_len_ptr = maybe_max_item_len_ptr ? static_cast(maybe_max_item_len_ptr->data_ptr()): nullptr; \ +params.additional_params.logits_soft_cap = logits_soft_cap; \ +params.additional_params.sm_scale = sm_scale; \ +params.additional_params.token_pos_in_items_len = token_pos_in_items_len; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, RaggedParams, PagedParams, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { using AttentionVariant = DefaultAttention; __VA_ARGS__();}) + +using namespace flashinfer; + +using DTypeQ = cutlass_dtype_t; +using DTypeKV = cutlass_dtype_t; +using DTypeO = cutlass_dtype_t; +using IdType = cutlass_dtype_t; + +constexpr int HEAD_DIM_QK = 128; +constexpr int HEAD_DIM_VO = 128; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto USE_SLIDING_WINDOW = false; + +struct RaggedParams { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + // The QKV matrices. + DTypeQ* q_ptr; + DTypeKV* k_ptr; + DTypeKV* v_ptr; + DTypeO* o_ptr; + float* lse_ptr; + + IdType* qo_tile_indices; + IdType* qo_indptr; + IdType* kv_indptr; + IdType* qo_lens; + IdType* kv_lens; + IdType* head_indices; + IdType* work_indptr; + IdType* batch_indices; + + struct AdditionalParams { + uint32_t* maybe_prefix_len_ptr; +uint16_t* maybe_token_pos_in_items_ptr; +uint16_t* maybe_max_item_len_ptr; +double logits_soft_cap; +double sm_scale; +int64_t token_pos_in_items_len; + + } additional_params; + + int64_t q_stride_n; + int64_t k_stride_n; + int64_t v_stride_n; + int64_t o_stride_n; + int64_t q_stride_h; + int64_t k_stride_h; + int64_t v_stride_h; + int64_t o_stride_h; + int64_t nnz_qo; + int64_t nnz_kv; + + int head_dim; + int num_qo_heads; + int num_kv_heads; + int group_size; + int window_left; + + bool causal; +}; + +struct PagedParams { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + // The QKV matrices. + DTypeQ* q_ptr; + DTypeKV* k_ptr; + DTypeKV* v_ptr; + DTypeO* o_ptr; + float* lse_ptr; + + IdType* qo_tile_indices; + IdType* qo_indptr; + IdType* kv_indptr; + IdType* kv_indices; + IdType* qo_lens; + IdType* kv_lens; + IdType* head_indices; + IdType* work_indptr; + IdType* batch_indices; + + struct AdditionalParams { + uint32_t* maybe_prefix_len_ptr; +uint16_t* maybe_token_pos_in_items_ptr; +uint16_t* maybe_max_item_len_ptr; +double logits_soft_cap; +double sm_scale; +int64_t token_pos_in_items_len; + + } additional_params; + + int64_t q_stride_n; + int64_t k_stride_n; + int64_t v_stride_n; + int64_t o_stride_n; + int64_t q_stride_h; + int64_t k_stride_h; + int64_t v_stride_h; + int64_t o_stride_h; + int64_t nnz_qo; + + int head_dim; + int num_qo_heads; + int num_kv_heads; + int group_size; + int page_size; + int window_left; + + bool causal; +}; + +#include \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_sm90_jit_pybind.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_sm90_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..b1edd582549e124c776f2cf9b4503f4bf60d16ae --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_sm90_jit_pybind.cu @@ -0,0 +1,46 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "batch_prefill_sm90_config.inc" +#include "pytorch_extension_utils.h" + +at::Tensor BatchPrefillWithKVCacheSM90Plan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, + int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, + int64_t head_dim_vo, bool causal); + +void BatchPrefillWithRaggedKVCacheSM90Run( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +void BatchPrefillWithPagedKVCacheSM90Run( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, + at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Batch-request prefill attention with KV-Cache plan + m.def("plan", BatchPrefillWithKVCacheSM90Plan); + // Batch-request prefill attention with KV-Cache operator + m.def("ragged_run", BatchPrefillWithRaggedKVCacheSM90Run); + // Batch-request prefill attention with KV-Cache operator + m.def("paged_run", BatchPrefillWithPagedKVCacheSM90Run); +} diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_0.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..857a0f5796850a0d4499248335a5c7f1e2dad25b --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_0.cu @@ -0,0 +1,24 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched + <192, + 128, + MaskMode::kNone, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultAttention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched + <192, + 128, + MaskMode::kNone, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultAttention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_1.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..2f298fe5b33cea4a229c62692d707ea5f8d6fce5 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_1.cu @@ -0,0 +1,24 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched + <192, + 128, + MaskMode::kCausal, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultAttention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched + <192, + 128, + MaskMode::kCausal, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultAttention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_2.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..d075a56a378fce5b3b4c3a14e75309f3db9e560f --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_2.cu @@ -0,0 +1,24 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched + <192, + 128, + MaskMode::kCustom, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultAttention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched + <192, + 128, + MaskMode::kCustom, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultAttention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_3.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..8f188b661bb83a6220f002a64d7f2bd0d91d482f --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_3.cu @@ -0,0 +1,24 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched + <192, + 128, + MaskMode::kMultiItemScoring, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultAttention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched + <192, + 128, + MaskMode::kMultiItemScoring, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultAttention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_0.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..a8ff81056119dcb0fe6a8e23661212bcbb9df29c --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_0.cu @@ -0,0 +1,24 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched + <192, + 128, + MaskMode::kNone, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultAttention>(RaggedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched + <192, + 128, + MaskMode::kNone, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultAttention>(RaggedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_1.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..298cab570997fa4520c656188979e48efa6bb65b --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_1.cu @@ -0,0 +1,24 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched + <192, + 128, + MaskMode::kCausal, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultAttention>(RaggedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched + <192, + 128, + MaskMode::kCausal, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultAttention>(RaggedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_2.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..40ef2acdb67e3f7748afb7c8dea728bbc407fd3d --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_2.cu @@ -0,0 +1,24 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched + <192, + 128, + MaskMode::kCustom, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultAttention>(RaggedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched + <192, + 128, + MaskMode::kCustom, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultAttention>(RaggedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_3.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..f50854c38ab7a25526c070ff8eb10550151ed35b --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_3.cu @@ -0,0 +1,24 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched + <192, + 128, + MaskMode::kMultiItemScoring, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultAttention>(RaggedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched + <192, + 128, + MaskMode::kMultiItemScoring, + /*USE_SLIDING_WINDOW=*/false, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultAttention>(RaggedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_sm90.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_sm90.cu new file mode 100644 index 0000000000000000000000000000000000000000..159ee3f1eda3f47769ec14f8fac16f35fa9bda05 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_sm90.cu @@ -0,0 +1,260 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include +#include + +#include "batch_prefill_sm90_config.inc" +#include "pytorch_conversion_utils.h" +#include "pytorch_extension_utils.h" + +namespace flashinfer { + +template +cudaError_t BatchPrefillWithRaggedKVCacheDispatched(Params& params, bool enable_pdl, + cudaStream_t stream); + +template +cudaError_t BatchPrefillWithPagedKVCacheDispatched(Params& params, bool enable_pdl, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +at::Tensor BatchPrefillWithKVCacheSM90Plan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, + int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, + int64_t head_dim_vo, bool causal) { + size_t float_workspace_size_in_bytes = + float_workspace_buffer.size(0) * float_workspace_buffer.element_size(); + size_t int_workspace_size_in_bytes = + int_workspace_buffer.size(0) * int_workspace_buffer.element_size(); + + flashinfer::PrefillPlanSM90Info plan_info; + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + cudaError_t status = + PrefillSM90Plan(float_workspace_buffer.data_ptr(), float_workspace_size_in_bytes, + int_workspace_buffer.data_ptr(), page_locked_int_workspace_buffer.data_ptr(), + int_workspace_size_in_bytes, plan_info, qo_indptr.data_ptr(), + kv_indptr.data_ptr(), kv_len_arr.data_ptr(), total_num_rows, + batch_size, num_qo_heads, num_kv_heads, head_dim_qk, head_dim_vo, page_size, + causal, enable_cuda_graph, /*sizeof_dtype_o=*/2, stream); + + TORCH_CHECK(status == cudaSuccess, + "PrefillSM90Plan failed with error: ", cudaGetErrorString(status)); + + return vec_to_tensor(plan_info.ToVector()); +} + +void BatchPrefillWithRaggedKVCacheSM90Run( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + PrefillPlanSM90Info plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == q.size(0), lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == q.size(1), lse.size(1), q.size(1)); + } + + void* float_buffer_ptr = float_workspace_buffer.data_ptr(); + void* int_buffer_ptr = int_workspace_buffer.data_ptr(); + + int64_t head_dim_qk = q.size(2); + int64_t head_dim_vo = v.size(2); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + QKVLayout kv_layout = static_cast(layout); + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + const MaskMode mask_mode = static_cast(mask_mode_code); + bool use_swa = window_left != -1; + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, USE_SLIDING_WINDOW, + USE_LOGITS_SOFT_CAP, AttentionVariant, RaggedParams, PagedParams, [&] { + RaggedParams params; + + params.q_ptr = static_cast(q.data_ptr()); + params.k_ptr = static_cast(k.data_ptr()); + params.v_ptr = static_cast(v.data_ptr()); + params.o_ptr = static_cast(o.data_ptr()); + params.lse_ptr = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.q_stride_n = q.stride(0); + params.q_stride_h = q.stride(1); + params.o_stride_n = o.stride(0); + params.o_stride_h = o.stride(1); + if (kv_layout == QKVLayout::kNHD) { + params.k_stride_n = k.stride(0); + params.k_stride_h = k.stride(1); + params.v_stride_n = v.stride(0); + params.v_stride_h = v.stride(1); + } else { + params.k_stride_h = k.stride(0); + params.k_stride_n = k.stride(1); + params.v_stride_h = v.stride(0); + params.v_stride_n = v.stride(1); + } + params.nnz_qo = q.size(0); + params.nnz_kv = k.size(0); + params.num_qo_heads = q.size(1); + params.num_kv_heads = k.size(1); + params.group_size = params.num_qo_heads / params.num_kv_heads; + params.window_left = window_left; + params.causal = mask_mode_code == 1; + params.qo_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_tile_indices_offset); + params.qo_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_indptr_offset); + params.kv_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_indptr_offset); + params.qo_lens = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_len_offset); + params.kv_lens = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_len_offset); + params.head_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.head_indices_offset); + params.work_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.work_indptr_offset); + params.batch_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.batch_indices_offset); + + ADDITIONAL_PARAMS_SETTER + + bool same_schedule_for_all_heads = plan_info.same_schedule_for_all_heads; + DISPATCH_BOOL(same_schedule_for_all_heads, SAME_SCHEDULER_FOR_ALL_HEADS, [&] { + cudaError_t status = BatchPrefillWithRaggedKVCacheDispatched< + HEAD_DIM_QK, HEAD_DIM_VO, MASK_MODE, USE_SLIDING_WINDOW, SAME_SCHEDULER_FOR_ALL_HEADS, + AttentionVariant>(params, enable_pdl, stream); + TORCH_CHECK(status == cudaSuccess, + "BatchPrefillWithRaggedKVCacheSM90Run failed with error: ", + cudaGetErrorString(status)); + return true; + }); + }); +} + +void BatchPrefillWithPagedKVCacheSM90Run( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, + at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + PrefillPlanSM90Info plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == q.size(0), lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == q.size(1), lse.size(1), q.size(1)); + } + QKVLayout kv_layout = static_cast(layout); + int64_t num_kv_heads, page_size; + int64_t head_dim_qk = q.size(2); + int64_t head_dim_vo = paged_v_cache.size(3); + if (kv_layout == QKVLayout::kHND) { + num_kv_heads = paged_k_cache.size(1); + page_size = paged_k_cache.size(2); + } else { + page_size = paged_k_cache.size(1); + num_kv_heads = paged_k_cache.size(2); + } + + void* float_buffer_ptr = float_workspace_buffer.data_ptr(); + void* int_buffer_ptr = int_workspace_buffer.data_ptr(); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = paged_k_cache.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + const MaskMode mask_mode = static_cast(mask_mode_code); + bool use_swa = window_left != -1; + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, USE_SLIDING_WINDOW, + USE_LOGITS_SOFT_CAP, AttentionVariant, RaggedParams, PagedParams, [&] { + PagedParams params; + + params.q_ptr = static_cast(q.data_ptr()); + params.k_ptr = static_cast(paged_k_cache.data_ptr()); + params.v_ptr = static_cast(paged_v_cache.data_ptr()); + params.o_ptr = static_cast(o.data_ptr()); + params.lse_ptr = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.q_stride_n = q.stride(0); + params.q_stride_h = q.stride(1); + params.o_stride_n = o.stride(0); + params.o_stride_h = o.stride(1); + if (kv_layout == QKVLayout::kNHD) { + // (num_pages, page_size, num_heads, head_dim) + params.k_stride_n = paged_k_cache.stride(1); + params.k_stride_h = paged_k_cache.stride(2); + params.v_stride_n = paged_v_cache.stride(1); + params.v_stride_h = paged_v_cache.stride(2); + } else { + // (num_pages, num_heads, page_size, head_dim) + params.k_stride_h = paged_k_cache.stride(1); + params.k_stride_n = paged_k_cache.stride(2); + params.v_stride_h = paged_v_cache.stride(1); + params.v_stride_n = paged_v_cache.stride(2); + } + params.nnz_qo = q.size(0); + params.num_qo_heads = q.size(1); + params.num_kv_heads = num_kv_heads; + params.group_size = params.num_qo_heads / num_kv_heads; + params.page_size = page_size; + params.window_left = window_left; + params.causal = mask_mode_code == 1; + params.qo_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_tile_indices_offset); + params.qo_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_indptr_offset); + params.kv_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_indptr_offset); + params.qo_lens = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_len_offset); + params.kv_lens = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_len_offset); + params.head_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.head_indices_offset); + params.work_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.work_indptr_offset); + params.batch_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.batch_indices_offset); + params.kv_indices = static_cast(paged_kv_indices.data_ptr()); + + ADDITIONAL_PARAMS_SETTER + + bool same_schedule_for_all_heads = plan_info.same_schedule_for_all_heads; + DISPATCH_BOOL(same_schedule_for_all_heads, SAME_SCHEDULER_FOR_ALL_HEADS, [&] { + cudaError_t status = BatchPrefillWithPagedKVCacheDispatched< + HEAD_DIM_QK, HEAD_DIM_VO, MASK_MODE, USE_SLIDING_WINDOW, SAME_SCHEDULER_FOR_ALL_HEADS, + AttentionVariant>(params, enable_pdl, stream); + TORCH_CHECK(status == cudaSuccess, + "BatchPrefillWithPagedKVCacheSM90Run failed with error: ", + cudaGetErrorString(status)); + return true; + }); + }); +} diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_sm90_config.inc b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_sm90_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..3603585516c0239187f303496b964492f561d2bf --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_sm90_config.inc @@ -0,0 +1,135 @@ +#pragma once +#include +#include +#include +#include +#include +#include "pytorch_extension_utils.h" + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_prefix_len_ptr, std::optional maybe_token_pos_in_items_ptr, std::optional maybe_max_item_len_ptr, double logits_soft_cap, double sm_scale, int64_t token_pos_in_items_len +#define ADDITIONAL_PARAMS_SETTER params.additional_params.maybe_prefix_len_ptr = maybe_prefix_len_ptr ? static_cast(maybe_prefix_len_ptr->data_ptr()): nullptr; \ +params.additional_params.maybe_token_pos_in_items_ptr = maybe_token_pos_in_items_ptr ? static_cast(maybe_token_pos_in_items_ptr->data_ptr()): nullptr; \ +params.additional_params.maybe_max_item_len_ptr = maybe_max_item_len_ptr ? static_cast(maybe_max_item_len_ptr->data_ptr()): nullptr; \ +params.additional_params.logits_soft_cap = logits_soft_cap; \ +params.additional_params.sm_scale = sm_scale; \ +params.additional_params.token_pos_in_items_len = token_pos_in_items_len; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, RaggedParams, PagedParams, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { using AttentionVariant = DefaultAttention; __VA_ARGS__();}) + +using namespace flashinfer; + +using DTypeQ = cutlass_dtype_t; +using DTypeKV = cutlass_dtype_t; +using DTypeO = cutlass_dtype_t; +using IdType = cutlass_dtype_t; + +constexpr int HEAD_DIM_QK = 192; +constexpr int HEAD_DIM_VO = 128; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto USE_SLIDING_WINDOW = false; + +struct RaggedParams { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + // The QKV matrices. + DTypeQ* q_ptr; + DTypeKV* k_ptr; + DTypeKV* v_ptr; + DTypeO* o_ptr; + float* lse_ptr; + + IdType* qo_tile_indices; + IdType* qo_indptr; + IdType* kv_indptr; + IdType* qo_lens; + IdType* kv_lens; + IdType* head_indices; + IdType* work_indptr; + IdType* batch_indices; + + struct AdditionalParams { + uint32_t* maybe_prefix_len_ptr; +uint16_t* maybe_token_pos_in_items_ptr; +uint16_t* maybe_max_item_len_ptr; +double logits_soft_cap; +double sm_scale; +int64_t token_pos_in_items_len; + + } additional_params; + + int64_t q_stride_n; + int64_t k_stride_n; + int64_t v_stride_n; + int64_t o_stride_n; + int64_t q_stride_h; + int64_t k_stride_h; + int64_t v_stride_h; + int64_t o_stride_h; + int64_t nnz_qo; + int64_t nnz_kv; + + int head_dim; + int num_qo_heads; + int num_kv_heads; + int group_size; + int window_left; + + bool causal; +}; + +struct PagedParams { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + // The QKV matrices. + DTypeQ* q_ptr; + DTypeKV* k_ptr; + DTypeKV* v_ptr; + DTypeO* o_ptr; + float* lse_ptr; + + IdType* qo_tile_indices; + IdType* qo_indptr; + IdType* kv_indptr; + IdType* kv_indices; + IdType* qo_lens; + IdType* kv_lens; + IdType* head_indices; + IdType* work_indptr; + IdType* batch_indices; + + struct AdditionalParams { + uint32_t* maybe_prefix_len_ptr; +uint16_t* maybe_token_pos_in_items_ptr; +uint16_t* maybe_max_item_len_ptr; +double logits_soft_cap; +double sm_scale; +int64_t token_pos_in_items_len; + + } additional_params; + + int64_t q_stride_n; + int64_t k_stride_n; + int64_t v_stride_n; + int64_t o_stride_n; + int64_t q_stride_h; + int64_t k_stride_h; + int64_t v_stride_h; + int64_t o_stride_h; + int64_t nnz_qo; + + int head_dim; + int num_qo_heads; + int num_kv_heads; + int group_size; + int page_size; + int window_left; + + bool causal; +}; + +#include \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_sm90_jit_pybind.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_sm90_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..b1edd582549e124c776f2cf9b4503f4bf60d16ae --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/batch_prefill_sm90_jit_pybind.cu @@ -0,0 +1,46 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "batch_prefill_sm90_config.inc" +#include "pytorch_extension_utils.h" + +at::Tensor BatchPrefillWithKVCacheSM90Plan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, + int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, + int64_t head_dim_vo, bool causal); + +void BatchPrefillWithRaggedKVCacheSM90Run( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +void BatchPrefillWithPagedKVCacheSM90Run( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, + at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Batch-request prefill attention with KV-Cache plan + m.def("plan", BatchPrefillWithKVCacheSM90Plan); + // Batch-request prefill attention with KV-Cache operator + m.def("ragged_run", BatchPrefillWithRaggedKVCacheSM90Run); + // Batch-request prefill attention with KV-Cache operator + m.def("paged_run", BatchPrefillWithPagedKVCacheSM90Run); +} diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill.cu new file mode 100644 index 0000000000000000000000000000000000000000..a51fc7f56a0af7482bd5d2a9d3f672854df69dfd --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill.cu @@ -0,0 +1,330 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include +#include + +#include "batch_prefill_config.inc" +#include "pytorch_conversion_utils.h" +#include "pytorch_extension_utils.h" + +namespace flashinfer { + +template +cudaError_t BatchPrefillWithPagedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v, + float* tmp_s, bool enable_pdl, + cudaStream_t stream); + +template +cudaError_t BatchPrefillWithRaggedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v, + float* tmp_s, bool enable_pdl, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +at::Tensor BatchPrefillWithKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, + int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, + int64_t head_dim_vo, bool causal) { + size_t float_workspace_size_in_bytes = + float_workspace_buffer.size(0) * float_workspace_buffer.element_size(); + size_t int_workspace_size_in_bytes = + int_workspace_buffer.size(0) * int_workspace_buffer.element_size(); + + PrefillPlanInfo plan_info; + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + cudaError_t status = PrefillPlan( + float_workspace_buffer.data_ptr(), float_workspace_size_in_bytes, + int_workspace_buffer.data_ptr(), page_locked_int_workspace_buffer.data_ptr(), + int_workspace_size_in_bytes, plan_info, qo_indptr.data_ptr(), + kv_indptr.data_ptr(), total_num_rows, batch_size, num_qo_heads, num_kv_heads, + head_dim_qk, head_dim_vo, page_size, enable_cuda_graph, /*sizeof_dtype_o=*/2, stream); + + TORCH_CHECK(status == cudaSuccess, + "Failed to plan prefill with error: ", cudaGetErrorString(status)); + + return vec_to_tensor(plan_info.ToVector()); +} + +void BatchPrefillWithRaggedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor k, at::Tensor v, + at::Tensor qo_indptr, at::Tensor kv_indptr, at::Tensor o, + std::optional maybe_lse, int64_t mask_mode_code, + int64_t layout, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + PrefillPlanInfo plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + QKVLayout kv_layout = static_cast(layout); + + int64_t num_qo_heads = q.size(1); + int64_t head_dim_qk = q.size(2); + int64_t num_kv_heads = (kv_layout == QKVLayout::kNHD) ? k.size(1) : k.size(0); + uint32_t q_stride_n = q.stride(0), q_stride_h = q.stride(1), k_stride_n, k_stride_h, v_stride_n, + v_stride_h; + if (kv_layout == QKVLayout::kNHD) { + k_stride_n = k.stride(0); + k_stride_h = k.stride(1); + v_stride_n = v.stride(0); + v_stride_h = v.stride(1); + } else { + k_stride_h = k.stride(0); + k_stride_n = k.stride(1); + v_stride_h = v.stride(0); + v_stride_n = v.stride(1); + } + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == q.size(0), lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == q.size(1), lse.size(1), q.size(1)); + } + + void* float_buffer_ptr = float_workspace_buffer.data_ptr(); + void* int_buffer_ptr = int_workspace_buffer.data_ptr(); + + const MaskMode mask_mode = static_cast(mask_mode_code); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, + RaggedParams, PagedParams, [&] { + RaggedParams params; + + params.q = static_cast(q.data_ptr()); + params.k = static_cast(k.data_ptr()); + params.v = static_cast(v.data_ptr()); + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.q_indptr = static_cast(qo_indptr.data_ptr()); + params.kv_indptr = static_cast(kv_indptr.data_ptr()); + params.num_qo_heads = num_qo_heads; + params.num_kv_heads = num_kv_heads; + params.group_size = uint_fastdiv(num_qo_heads / num_kv_heads); + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.k_stride_n = k_stride_n; + params.k_stride_h = k_stride_h; + params.v_stride_n = v_stride_n; + params.v_stride_h = v_stride_h; + params.window_left = window_left; + + params.request_indices = nullptr; + params.qo_tile_indices = nullptr; + params.kv_tile_indices = nullptr; + params.merge_indptr = nullptr; + params.o_indptr = nullptr; + params.kv_chunk_size_ptr = nullptr; + params.block_valid_mask = nullptr; + params.total_num_rows = nullptr; + params.max_total_num_rows = 0; + params.padded_batch_size = 0; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + DTypeO* tmp_v = nullptr; + float* tmp_s = nullptr; + + params.request_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.request_indices_offset); + params.qo_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_tile_indices_offset); + params.kv_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_tile_indices_offset); + params.o_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.o_indptr_offset); + params.kv_chunk_size_ptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_chunk_size_ptr_offset); + if (plan_info.split_kv) { + params.merge_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.merge_indptr_offset); + tmp_v = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.v_offset); + tmp_s = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.s_offset); + if (plan_info.enable_cuda_graph) { + params.block_valid_mask = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.block_valid_mask_offset); + } + } + params.padded_batch_size = plan_info.padded_batch_size; + params.max_total_num_rows = plan_info.total_num_rows; + if (plan_info.enable_cuda_graph) { + params.total_num_rows = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.total_num_rows_offset); + } + + cudaError_t status = cudaSuccess; + + DISPATCH_CTA_TILE_Q(plan_info.cta_tile_q, CTA_TILE_Q, { + status = flashinfer::BatchPrefillWithRaggedKVCacheDispatched< + CTA_TILE_Q, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + /*use_fp16_qk_reduction=*/USE_FP16_QK_REDUCTION, MASK_MODE, AttentionVariant, + RaggedParams>(params, tmp_v, tmp_s, enable_pdl, stream); + }); + + TORCH_CHECK(status == cudaSuccess, "BatchPrefillWithRaggedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); +} + +void BatchPrefillWithPagedKVCacheRun( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, + at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + PrefillPlanInfo plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + QKVLayout kv_layout = static_cast(layout); + auto device = q.device(); + int64_t batch_size = paged_kv_indptr.size(0) - 1; + int64_t num_qo_heads = q.size(1); + int64_t num_kv_heads, page_size; + uint32_t head_dim_qk = q.size(2); + if (kv_layout == QKVLayout::kHND) { + num_kv_heads = paged_k_cache.size(1); + page_size = paged_k_cache.size(2); + } else { + page_size = paged_k_cache.size(1); + num_kv_heads = paged_k_cache.size(2); + } + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == q.size(0), lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == q.size(1), lse.size(1), q.size(1)); + } + + void* float_buffer_ptr = static_cast(float_workspace_buffer.data_ptr()); + void* int_buffer_ptr = static_cast(int_workspace_buffer.data_ptr()); + + const MaskMode mask_mode = static_cast(mask_mode_code); + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = paged_k_cache.scalar_type(); + + // get q_stride_n and q_stride_h + const auto q_stride_n = q.stride(0); + const auto q_stride_h = q.stride(1); + + // get kv_cache_strides + const int64_t* kv_cache_strides = nullptr; + auto k_strides = paged_k_cache.strides(); + auto v_strides = paged_v_cache.strides(); + TORCH_CHECK(k_strides == v_strides, "k/v strides must be identical"); + kv_cache_strides = k_strides.data(); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, + RaggedParams, PagedParams, [&] { + PagedParams params; + + params.q = static_cast(q.data_ptr()); + paged_kv_t paged_kv( + num_kv_heads, page_size, HEAD_DIM_VO, batch_size, kv_layout, + static_cast(paged_k_cache.data_ptr()), + static_cast(paged_v_cache.data_ptr()), kv_cache_strides, + static_cast(paged_kv_indices.data_ptr()), + static_cast(paged_kv_indptr.data_ptr()), + static_cast(paged_kv_last_page_len.data_ptr())); + params.paged_kv = paged_kv; + params.q_indptr = static_cast(qo_indptr.data_ptr()); + params.o = static_cast(o.data_ptr()); + + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.num_qo_heads = num_qo_heads; + params.group_size = uint_fastdiv(num_qo_heads / paged_kv.num_heads); + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.window_left = window_left; + + params.request_indices = nullptr; + params.qo_tile_indices = nullptr; + params.kv_tile_indices = nullptr; + params.merge_indptr = nullptr; + params.o_indptr = nullptr; + params.kv_chunk_size_ptr = nullptr; + params.block_valid_mask = nullptr; + params.total_num_rows = nullptr; + params.max_total_num_rows = 0; + params.padded_batch_size = 0; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + DTypeO* tmp_v = nullptr; + float* tmp_s = nullptr; + + params.request_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.request_indices_offset); + params.qo_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_tile_indices_offset); + params.kv_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_tile_indices_offset); + params.o_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.o_indptr_offset); + params.kv_chunk_size_ptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_chunk_size_ptr_offset); + if (plan_info.split_kv) { + params.merge_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.merge_indptr_offset); + tmp_v = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.v_offset); + tmp_s = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.s_offset); + if (plan_info.enable_cuda_graph) { + params.block_valid_mask = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.block_valid_mask_offset); + } + } + params.padded_batch_size = plan_info.padded_batch_size; + params.max_total_num_rows = plan_info.total_num_rows; + if (plan_info.enable_cuda_graph) { + params.total_num_rows = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.total_num_rows_offset); + } + + cudaError_t status = cudaSuccess; + + DISPATCH_CTA_TILE_Q(plan_info.cta_tile_q, CTA_TILE_Q, { + status = flashinfer::BatchPrefillWithPagedKVCacheDispatched< + CTA_TILE_Q, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + /*use_fp16_qk_reduction=*/USE_FP16_QK_REDUCTION, MASK_MODE, AttentionVariant, + PagedParams>(params, tmp_v, tmp_s, enable_pdl, stream); + }); + + TORCH_CHECK(status == cudaSuccess, "BatchPrefillWithPagedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); +} diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_config.inc b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..6b262acf92536a8ac576f651305ae9b688dd6c4e --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_config.inc @@ -0,0 +1,153 @@ +#pragma once +#include +#include +#include +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_custom_mask, std::optional maybe_mask_indptr, std::optional maybe_alibi_slopes, std::optional maybe_prefix_len_ptr, std::optional maybe_token_pos_in_items_ptr, std::optional maybe_max_item_len_ptr, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta, int64_t token_pos_in_items_len +#define ADDITIONAL_PARAMS_SETTER params.maybe_custom_mask = maybe_custom_mask ? static_cast(maybe_custom_mask->data_ptr()): nullptr; \ +params.maybe_mask_indptr = maybe_mask_indptr ? static_cast(maybe_mask_indptr->data_ptr()): nullptr; \ +params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.maybe_prefix_len_ptr = maybe_prefix_len_ptr ? static_cast(maybe_prefix_len_ptr->data_ptr()): nullptr; \ +params.maybe_token_pos_in_items_ptr = maybe_token_pos_in_items_ptr ? static_cast(maybe_token_pos_in_items_ptr->data_ptr()): nullptr; \ +params.maybe_max_item_len_ptr = maybe_max_item_len_ptr ? static_cast(maybe_max_item_len_ptr->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; \ +params.token_pos_in_items_len = token_pos_in_items_len; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, RaggedParams, PagedParams, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { \ + constexpr auto use_custom_mask = MASK_MODE == MaskMode::kCustom; \ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ + }) + +using namespace flashinfer; + +using DTypeQ = half; +using DTypeKV = half; +using DTypeO = half; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 256; +constexpr int HEAD_DIM_VO = 256; +constexpr bool USE_FP16_QK_REDUCTION = false; +constexpr auto USE_LOGITS_SOFT_CAP = true; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = true; + + +struct RaggedParams { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + DTypeQ* q; + DTypeKV* k; + DTypeKV* v; + IdType* q_indptr; + IdType* kv_indptr; + DTypeO* o; + float* lse; + uint_fastdiv group_size; + + uint8_t* maybe_custom_mask; +int32_t* maybe_mask_indptr; +float* maybe_alibi_slopes; +uint32_t* maybe_prefix_len_ptr; +uint16_t* maybe_token_pos_in_items_ptr; +uint16_t* maybe_max_item_len_ptr; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; +int64_t token_pos_in_items_len; + + uint32_t num_qo_heads; + uint32_t num_kv_heads; + uint32_t q_stride_n; + uint32_t q_stride_h; + uint32_t k_stride_n; + uint32_t k_stride_h; + uint32_t v_stride_n; + uint32_t v_stride_h; + int32_t window_left; + + IdType* request_indices; + IdType* qo_tile_indices; + IdType* kv_tile_indices; + IdType* merge_indptr; + IdType* o_indptr; + IdType* kv_chunk_size_ptr; + bool* block_valid_mask; + uint32_t max_total_num_rows; + uint32_t* total_num_rows; + uint32_t padded_batch_size; + bool partition_kv; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { + return q_indptr[batch_idx + 1] - q_indptr[batch_idx]; + } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return kv_indptr[batch_idx + 1] - kv_indptr[batch_idx]; + } +}; + +struct PagedParams { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + DTypeQ* q; + paged_kv_t paged_kv; + IdType* q_indptr; + DTypeO* o; + float* lse; + uint_fastdiv group_size; + + uint8_t* maybe_custom_mask; +int32_t* maybe_mask_indptr; +float* maybe_alibi_slopes; +uint32_t* maybe_prefix_len_ptr; +uint16_t* maybe_token_pos_in_items_ptr; +uint16_t* maybe_max_item_len_ptr; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; +int64_t token_pos_in_items_len; + + uint32_t num_qo_heads; + IdType q_stride_n; + IdType q_stride_h; + int32_t window_left; + + IdType* request_indices; + IdType* qo_tile_indices; + IdType* kv_tile_indices; + IdType* merge_indptr; + IdType* o_indptr; + bool* block_valid_mask; + IdType* kv_chunk_size_ptr; + uint32_t max_total_num_rows; + uint32_t* total_num_rows; + uint32_t padded_batch_size; + bool partition_kv; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { + return q_indptr[batch_idx + 1] - q_indptr[batch_idx]; + } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return paged_kv.get_length(batch_idx); + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_jit_pybind.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..5421ab1cf910fdfbc867c182c8e22f23e44749ae --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_jit_pybind.cu @@ -0,0 +1,48 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "batch_prefill_config.inc" +#include "pytorch_extension_utils.h" + +at::Tensor BatchPrefillWithKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, + int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, + int64_t head_dim_vo, bool causal); + +void BatchPrefillWithRaggedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor k, at::Tensor v, + at::Tensor qo_indptr, at::Tensor kv_indptr, at::Tensor o, + std::optional maybe_lse, int64_t mask_mode_code, + int64_t layout, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +void BatchPrefillWithPagedKVCacheRun( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, + at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Batch-request prefill attention with KV-Cache plan + m.def("plan", BatchPrefillWithKVCachePlan); + // Batch-request prefill attention with KV-Cache operator + m.def("ragged_run", BatchPrefillWithRaggedKVCacheRun); + // Batch-request prefill attention with KV-Cache operator + m.def("paged_run", BatchPrefillWithPagedKVCacheRun); +} diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_paged_kernel_mask_0.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_paged_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..9ed3e16f02b81516d03b8e2b790df2742eb1430d --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_paged_kernel_mask_0.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kNone == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 256, 256, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 256, 256, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 256, 256, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_paged_kernel_mask_1.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_paged_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..37d4bfe8cdbda0809bc8b11730692da3632a1a44 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_paged_kernel_mask_1.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCausal == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 256, 256, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 256, 256, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 256, 256, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_paged_kernel_mask_2.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_paged_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..6d9b56ea582f9197cf55df1a48e80b7efce8273e --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_paged_kernel_mask_2.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCustom == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 256, 256, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 256, 256, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 256, 256, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_paged_kernel_mask_3.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_paged_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..4e67f0a63d0ba0aa6c2ef1be3f9f7e8e5bbac259 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_paged_kernel_mask_3.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kMultiItemScoring == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 256, 256, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 256, 256, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 256, 256, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_ragged_kernel_mask_0.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_ragged_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..4e4965782d385f94d63c616c3be1fdeef33474cb --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_ragged_kernel_mask_0.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kNone == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 256, 256, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 256, 256, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 256, 256, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_ragged_kernel_mask_1.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_ragged_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..eff49053b01f63913678acf1f2b86903f358de93 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_ragged_kernel_mask_1.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCausal == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 256, 256, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 256, 256, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 256, 256, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_ragged_kernel_mask_2.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_ragged_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..0c827cffe82db344b2afee6e33487871112ac672 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_ragged_kernel_mask_2.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCustom == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 256, 256, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 256, 256, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 256, 256, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_ragged_kernel_mask_3.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_ragged_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..5d7767dc6d8145fe1622713c8b1acbb1a61e4c8c --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/batch_prefill_ragged_kernel_mask_3.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kMultiItemScoring == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 256, 256, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 256, 256, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 256, 256, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_0.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..700eb2e367c50b3f4e1bf963849d60e099c03a99 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_0.cu @@ -0,0 +1,24 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched + <256, + 256, + MaskMode::kNone, + /*USE_SLIDING_WINDOW=*/true, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultAttention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched + <256, + 256, + MaskMode::kNone, + /*USE_SLIDING_WINDOW=*/true, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultAttention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_1.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..ae7ae1a87dc15bcbb01569c0ad7fa25ef0f3e8f7 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_1.cu @@ -0,0 +1,24 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched + <256, + 256, + MaskMode::kCausal, + /*USE_SLIDING_WINDOW=*/true, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultAttention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched + <256, + 256, + MaskMode::kCausal, + /*USE_SLIDING_WINDOW=*/true, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultAttention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_2.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..d2a680686e6ea2a4611481277d89a2fc3608d9f4 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_2.cu @@ -0,0 +1,24 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched + <256, + 256, + MaskMode::kCustom, + /*USE_SLIDING_WINDOW=*/true, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultAttention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched + <256, + 256, + MaskMode::kCustom, + /*USE_SLIDING_WINDOW=*/true, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultAttention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_3.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..4a0f46b6c308a6fd4dcf12ae2d4be11fa5bd89eb --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_paged_sm90_kernel_mask_3.cu @@ -0,0 +1,24 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched + <256, + 256, + MaskMode::kMultiItemScoring, + /*USE_SLIDING_WINDOW=*/true, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultAttention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched + <256, + 256, + MaskMode::kMultiItemScoring, + /*USE_SLIDING_WINDOW=*/true, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultAttention, PagedParams>(PagedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_0.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..afd17ddcf65342899385c0094d7187f93b92756d --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_0.cu @@ -0,0 +1,24 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched + <256, + 256, + MaskMode::kNone, + /*USE_SLIDING_WINDOW=*/true, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultAttention>(RaggedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched + <256, + 256, + MaskMode::kNone, + /*USE_SLIDING_WINDOW=*/true, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultAttention>(RaggedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_1.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..c901f44d76128501cc772b356f5217211199bb50 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_1.cu @@ -0,0 +1,24 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched + <256, + 256, + MaskMode::kCausal, + /*USE_SLIDING_WINDOW=*/true, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultAttention>(RaggedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched + <256, + 256, + MaskMode::kCausal, + /*USE_SLIDING_WINDOW=*/true, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultAttention>(RaggedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_2.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..00b399ef1c529610375703d79f261333f1c1c386 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_2.cu @@ -0,0 +1,24 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched + <256, + 256, + MaskMode::kCustom, + /*USE_SLIDING_WINDOW=*/true, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultAttention>(RaggedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched + <256, + 256, + MaskMode::kCustom, + /*USE_SLIDING_WINDOW=*/true, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultAttention>(RaggedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_3.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..d9da788a703f3c97c596380ea9eccb4427ed228a --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_ragged_sm90_kernel_mask_3.cu @@ -0,0 +1,24 @@ +#include +#include "batch_prefill_sm90_config.inc" + +namespace flashinfer { + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched + <256, + 256, + MaskMode::kMultiItemScoring, + /*USE_SLIDING_WINDOW=*/true, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/true, + DefaultAttention>(RaggedParams& params, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched + <256, + 256, + MaskMode::kMultiItemScoring, + /*USE_SLIDING_WINDOW=*/true, + /*SAME_SCHEDULER_FOR_ALL_HEADS=*/false, + DefaultAttention>(RaggedParams& params, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_sm90.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_sm90.cu new file mode 100644 index 0000000000000000000000000000000000000000..159ee3f1eda3f47769ec14f8fac16f35fa9bda05 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_sm90.cu @@ -0,0 +1,260 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include +#include + +#include "batch_prefill_sm90_config.inc" +#include "pytorch_conversion_utils.h" +#include "pytorch_extension_utils.h" + +namespace flashinfer { + +template +cudaError_t BatchPrefillWithRaggedKVCacheDispatched(Params& params, bool enable_pdl, + cudaStream_t stream); + +template +cudaError_t BatchPrefillWithPagedKVCacheDispatched(Params& params, bool enable_pdl, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +at::Tensor BatchPrefillWithKVCacheSM90Plan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, + int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, + int64_t head_dim_vo, bool causal) { + size_t float_workspace_size_in_bytes = + float_workspace_buffer.size(0) * float_workspace_buffer.element_size(); + size_t int_workspace_size_in_bytes = + int_workspace_buffer.size(0) * int_workspace_buffer.element_size(); + + flashinfer::PrefillPlanSM90Info plan_info; + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + cudaError_t status = + PrefillSM90Plan(float_workspace_buffer.data_ptr(), float_workspace_size_in_bytes, + int_workspace_buffer.data_ptr(), page_locked_int_workspace_buffer.data_ptr(), + int_workspace_size_in_bytes, plan_info, qo_indptr.data_ptr(), + kv_indptr.data_ptr(), kv_len_arr.data_ptr(), total_num_rows, + batch_size, num_qo_heads, num_kv_heads, head_dim_qk, head_dim_vo, page_size, + causal, enable_cuda_graph, /*sizeof_dtype_o=*/2, stream); + + TORCH_CHECK(status == cudaSuccess, + "PrefillSM90Plan failed with error: ", cudaGetErrorString(status)); + + return vec_to_tensor(plan_info.ToVector()); +} + +void BatchPrefillWithRaggedKVCacheSM90Run( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + PrefillPlanSM90Info plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == q.size(0), lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == q.size(1), lse.size(1), q.size(1)); + } + + void* float_buffer_ptr = float_workspace_buffer.data_ptr(); + void* int_buffer_ptr = int_workspace_buffer.data_ptr(); + + int64_t head_dim_qk = q.size(2); + int64_t head_dim_vo = v.size(2); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + QKVLayout kv_layout = static_cast(layout); + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + const MaskMode mask_mode = static_cast(mask_mode_code); + bool use_swa = window_left != -1; + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, USE_SLIDING_WINDOW, + USE_LOGITS_SOFT_CAP, AttentionVariant, RaggedParams, PagedParams, [&] { + RaggedParams params; + + params.q_ptr = static_cast(q.data_ptr()); + params.k_ptr = static_cast(k.data_ptr()); + params.v_ptr = static_cast(v.data_ptr()); + params.o_ptr = static_cast(o.data_ptr()); + params.lse_ptr = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.q_stride_n = q.stride(0); + params.q_stride_h = q.stride(1); + params.o_stride_n = o.stride(0); + params.o_stride_h = o.stride(1); + if (kv_layout == QKVLayout::kNHD) { + params.k_stride_n = k.stride(0); + params.k_stride_h = k.stride(1); + params.v_stride_n = v.stride(0); + params.v_stride_h = v.stride(1); + } else { + params.k_stride_h = k.stride(0); + params.k_stride_n = k.stride(1); + params.v_stride_h = v.stride(0); + params.v_stride_n = v.stride(1); + } + params.nnz_qo = q.size(0); + params.nnz_kv = k.size(0); + params.num_qo_heads = q.size(1); + params.num_kv_heads = k.size(1); + params.group_size = params.num_qo_heads / params.num_kv_heads; + params.window_left = window_left; + params.causal = mask_mode_code == 1; + params.qo_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_tile_indices_offset); + params.qo_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_indptr_offset); + params.kv_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_indptr_offset); + params.qo_lens = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_len_offset); + params.kv_lens = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_len_offset); + params.head_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.head_indices_offset); + params.work_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.work_indptr_offset); + params.batch_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.batch_indices_offset); + + ADDITIONAL_PARAMS_SETTER + + bool same_schedule_for_all_heads = plan_info.same_schedule_for_all_heads; + DISPATCH_BOOL(same_schedule_for_all_heads, SAME_SCHEDULER_FOR_ALL_HEADS, [&] { + cudaError_t status = BatchPrefillWithRaggedKVCacheDispatched< + HEAD_DIM_QK, HEAD_DIM_VO, MASK_MODE, USE_SLIDING_WINDOW, SAME_SCHEDULER_FOR_ALL_HEADS, + AttentionVariant>(params, enable_pdl, stream); + TORCH_CHECK(status == cudaSuccess, + "BatchPrefillWithRaggedKVCacheSM90Run failed with error: ", + cudaGetErrorString(status)); + return true; + }); + }); +} + +void BatchPrefillWithPagedKVCacheSM90Run( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, + at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + PrefillPlanSM90Info plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == q.size(0), lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == q.size(1), lse.size(1), q.size(1)); + } + QKVLayout kv_layout = static_cast(layout); + int64_t num_kv_heads, page_size; + int64_t head_dim_qk = q.size(2); + int64_t head_dim_vo = paged_v_cache.size(3); + if (kv_layout == QKVLayout::kHND) { + num_kv_heads = paged_k_cache.size(1); + page_size = paged_k_cache.size(2); + } else { + page_size = paged_k_cache.size(1); + num_kv_heads = paged_k_cache.size(2); + } + + void* float_buffer_ptr = float_workspace_buffer.data_ptr(); + void* int_buffer_ptr = int_workspace_buffer.data_ptr(); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = paged_k_cache.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + const MaskMode mask_mode = static_cast(mask_mode_code); + bool use_swa = window_left != -1; + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, USE_SLIDING_WINDOW, + USE_LOGITS_SOFT_CAP, AttentionVariant, RaggedParams, PagedParams, [&] { + PagedParams params; + + params.q_ptr = static_cast(q.data_ptr()); + params.k_ptr = static_cast(paged_k_cache.data_ptr()); + params.v_ptr = static_cast(paged_v_cache.data_ptr()); + params.o_ptr = static_cast(o.data_ptr()); + params.lse_ptr = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.q_stride_n = q.stride(0); + params.q_stride_h = q.stride(1); + params.o_stride_n = o.stride(0); + params.o_stride_h = o.stride(1); + if (kv_layout == QKVLayout::kNHD) { + // (num_pages, page_size, num_heads, head_dim) + params.k_stride_n = paged_k_cache.stride(1); + params.k_stride_h = paged_k_cache.stride(2); + params.v_stride_n = paged_v_cache.stride(1); + params.v_stride_h = paged_v_cache.stride(2); + } else { + // (num_pages, num_heads, page_size, head_dim) + params.k_stride_h = paged_k_cache.stride(1); + params.k_stride_n = paged_k_cache.stride(2); + params.v_stride_h = paged_v_cache.stride(1); + params.v_stride_n = paged_v_cache.stride(2); + } + params.nnz_qo = q.size(0); + params.num_qo_heads = q.size(1); + params.num_kv_heads = num_kv_heads; + params.group_size = params.num_qo_heads / num_kv_heads; + params.page_size = page_size; + params.window_left = window_left; + params.causal = mask_mode_code == 1; + params.qo_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_tile_indices_offset); + params.qo_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_indptr_offset); + params.kv_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_indptr_offset); + params.qo_lens = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_len_offset); + params.kv_lens = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_len_offset); + params.head_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.head_indices_offset); + params.work_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.work_indptr_offset); + params.batch_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.batch_indices_offset); + params.kv_indices = static_cast(paged_kv_indices.data_ptr()); + + ADDITIONAL_PARAMS_SETTER + + bool same_schedule_for_all_heads = plan_info.same_schedule_for_all_heads; + DISPATCH_BOOL(same_schedule_for_all_heads, SAME_SCHEDULER_FOR_ALL_HEADS, [&] { + cudaError_t status = BatchPrefillWithPagedKVCacheDispatched< + HEAD_DIM_QK, HEAD_DIM_VO, MASK_MODE, USE_SLIDING_WINDOW, SAME_SCHEDULER_FOR_ALL_HEADS, + AttentionVariant>(params, enable_pdl, stream); + TORCH_CHECK(status == cudaSuccess, + "BatchPrefillWithPagedKVCacheSM90Run failed with error: ", + cudaGetErrorString(status)); + return true; + }); + }); +} diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_sm90_config.inc b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_sm90_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..646938fa30bb054aa552741c33cb15948d4ed357 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_sm90_config.inc @@ -0,0 +1,135 @@ +#pragma once +#include +#include +#include +#include +#include +#include "pytorch_extension_utils.h" + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_prefix_len_ptr, std::optional maybe_token_pos_in_items_ptr, std::optional maybe_max_item_len_ptr, double logits_soft_cap, double sm_scale, int64_t token_pos_in_items_len +#define ADDITIONAL_PARAMS_SETTER params.additional_params.maybe_prefix_len_ptr = maybe_prefix_len_ptr ? static_cast(maybe_prefix_len_ptr->data_ptr()): nullptr; \ +params.additional_params.maybe_token_pos_in_items_ptr = maybe_token_pos_in_items_ptr ? static_cast(maybe_token_pos_in_items_ptr->data_ptr()): nullptr; \ +params.additional_params.maybe_max_item_len_ptr = maybe_max_item_len_ptr ? static_cast(maybe_max_item_len_ptr->data_ptr()): nullptr; \ +params.additional_params.logits_soft_cap = logits_soft_cap; \ +params.additional_params.sm_scale = sm_scale; \ +params.additional_params.token_pos_in_items_len = token_pos_in_items_len; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, RaggedParams, PagedParams, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { using AttentionVariant = DefaultAttention; __VA_ARGS__();}) + +using namespace flashinfer; + +using DTypeQ = cutlass_dtype_t; +using DTypeKV = cutlass_dtype_t; +using DTypeO = cutlass_dtype_t; +using IdType = cutlass_dtype_t; + +constexpr int HEAD_DIM_QK = 256; +constexpr int HEAD_DIM_VO = 256; +constexpr auto USE_LOGITS_SOFT_CAP = true; +constexpr auto USE_SLIDING_WINDOW = true; + +struct RaggedParams { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + // The QKV matrices. + DTypeQ* q_ptr; + DTypeKV* k_ptr; + DTypeKV* v_ptr; + DTypeO* o_ptr; + float* lse_ptr; + + IdType* qo_tile_indices; + IdType* qo_indptr; + IdType* kv_indptr; + IdType* qo_lens; + IdType* kv_lens; + IdType* head_indices; + IdType* work_indptr; + IdType* batch_indices; + + struct AdditionalParams { + uint32_t* maybe_prefix_len_ptr; +uint16_t* maybe_token_pos_in_items_ptr; +uint16_t* maybe_max_item_len_ptr; +double logits_soft_cap; +double sm_scale; +int64_t token_pos_in_items_len; + + } additional_params; + + int64_t q_stride_n; + int64_t k_stride_n; + int64_t v_stride_n; + int64_t o_stride_n; + int64_t q_stride_h; + int64_t k_stride_h; + int64_t v_stride_h; + int64_t o_stride_h; + int64_t nnz_qo; + int64_t nnz_kv; + + int head_dim; + int num_qo_heads; + int num_kv_heads; + int group_size; + int window_left; + + bool causal; +}; + +struct PagedParams { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + // The QKV matrices. + DTypeQ* q_ptr; + DTypeKV* k_ptr; + DTypeKV* v_ptr; + DTypeO* o_ptr; + float* lse_ptr; + + IdType* qo_tile_indices; + IdType* qo_indptr; + IdType* kv_indptr; + IdType* kv_indices; + IdType* qo_lens; + IdType* kv_lens; + IdType* head_indices; + IdType* work_indptr; + IdType* batch_indices; + + struct AdditionalParams { + uint32_t* maybe_prefix_len_ptr; +uint16_t* maybe_token_pos_in_items_ptr; +uint16_t* maybe_max_item_len_ptr; +double logits_soft_cap; +double sm_scale; +int64_t token_pos_in_items_len; + + } additional_params; + + int64_t q_stride_n; + int64_t k_stride_n; + int64_t v_stride_n; + int64_t o_stride_n; + int64_t q_stride_h; + int64_t k_stride_h; + int64_t v_stride_h; + int64_t o_stride_h; + int64_t nnz_qo; + + int head_dim; + int num_qo_heads; + int num_kv_heads; + int group_size; + int page_size; + int window_left; + + bool causal; +}; + +#include \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_sm90_jit_pybind.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_sm90_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..b1edd582549e124c776f2cf9b4503f4bf60d16ae --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/batch_prefill_sm90_jit_pybind.cu @@ -0,0 +1,46 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "batch_prefill_sm90_config.inc" +#include "pytorch_extension_utils.h" + +at::Tensor BatchPrefillWithKVCacheSM90Plan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, + int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, + int64_t head_dim_vo, bool causal); + +void BatchPrefillWithRaggedKVCacheSM90Run( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +void BatchPrefillWithPagedKVCacheSM90Run( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, + at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Batch-request prefill attention with KV-Cache plan + m.def("plan", BatchPrefillWithKVCacheSM90Plan); + // Batch-request prefill attention with KV-Cache operator + m.def("ragged_run", BatchPrefillWithRaggedKVCacheSM90Run); + // Batch-request prefill attention with KV-Cache operator + m.def("paged_run", BatchPrefillWithPagedKVCacheSM90Run); +} diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill.cu new file mode 100644 index 0000000000000000000000000000000000000000..a51fc7f56a0af7482bd5d2a9d3f672854df69dfd --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill.cu @@ -0,0 +1,330 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include +#include + +#include "batch_prefill_config.inc" +#include "pytorch_conversion_utils.h" +#include "pytorch_extension_utils.h" + +namespace flashinfer { + +template +cudaError_t BatchPrefillWithPagedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v, + float* tmp_s, bool enable_pdl, + cudaStream_t stream); + +template +cudaError_t BatchPrefillWithRaggedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v, + float* tmp_s, bool enable_pdl, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +at::Tensor BatchPrefillWithKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, + int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, + int64_t head_dim_vo, bool causal) { + size_t float_workspace_size_in_bytes = + float_workspace_buffer.size(0) * float_workspace_buffer.element_size(); + size_t int_workspace_size_in_bytes = + int_workspace_buffer.size(0) * int_workspace_buffer.element_size(); + + PrefillPlanInfo plan_info; + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + cudaError_t status = PrefillPlan( + float_workspace_buffer.data_ptr(), float_workspace_size_in_bytes, + int_workspace_buffer.data_ptr(), page_locked_int_workspace_buffer.data_ptr(), + int_workspace_size_in_bytes, plan_info, qo_indptr.data_ptr(), + kv_indptr.data_ptr(), total_num_rows, batch_size, num_qo_heads, num_kv_heads, + head_dim_qk, head_dim_vo, page_size, enable_cuda_graph, /*sizeof_dtype_o=*/2, stream); + + TORCH_CHECK(status == cudaSuccess, + "Failed to plan prefill with error: ", cudaGetErrorString(status)); + + return vec_to_tensor(plan_info.ToVector()); +} + +void BatchPrefillWithRaggedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor k, at::Tensor v, + at::Tensor qo_indptr, at::Tensor kv_indptr, at::Tensor o, + std::optional maybe_lse, int64_t mask_mode_code, + int64_t layout, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + PrefillPlanInfo plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + QKVLayout kv_layout = static_cast(layout); + + int64_t num_qo_heads = q.size(1); + int64_t head_dim_qk = q.size(2); + int64_t num_kv_heads = (kv_layout == QKVLayout::kNHD) ? k.size(1) : k.size(0); + uint32_t q_stride_n = q.stride(0), q_stride_h = q.stride(1), k_stride_n, k_stride_h, v_stride_n, + v_stride_h; + if (kv_layout == QKVLayout::kNHD) { + k_stride_n = k.stride(0); + k_stride_h = k.stride(1); + v_stride_n = v.stride(0); + v_stride_h = v.stride(1); + } else { + k_stride_h = k.stride(0); + k_stride_n = k.stride(1); + v_stride_h = v.stride(0); + v_stride_n = v.stride(1); + } + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == q.size(0), lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == q.size(1), lse.size(1), q.size(1)); + } + + void* float_buffer_ptr = float_workspace_buffer.data_ptr(); + void* int_buffer_ptr = int_workspace_buffer.data_ptr(); + + const MaskMode mask_mode = static_cast(mask_mode_code); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, + RaggedParams, PagedParams, [&] { + RaggedParams params; + + params.q = static_cast(q.data_ptr()); + params.k = static_cast(k.data_ptr()); + params.v = static_cast(v.data_ptr()); + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.q_indptr = static_cast(qo_indptr.data_ptr()); + params.kv_indptr = static_cast(kv_indptr.data_ptr()); + params.num_qo_heads = num_qo_heads; + params.num_kv_heads = num_kv_heads; + params.group_size = uint_fastdiv(num_qo_heads / num_kv_heads); + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.k_stride_n = k_stride_n; + params.k_stride_h = k_stride_h; + params.v_stride_n = v_stride_n; + params.v_stride_h = v_stride_h; + params.window_left = window_left; + + params.request_indices = nullptr; + params.qo_tile_indices = nullptr; + params.kv_tile_indices = nullptr; + params.merge_indptr = nullptr; + params.o_indptr = nullptr; + params.kv_chunk_size_ptr = nullptr; + params.block_valid_mask = nullptr; + params.total_num_rows = nullptr; + params.max_total_num_rows = 0; + params.padded_batch_size = 0; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + DTypeO* tmp_v = nullptr; + float* tmp_s = nullptr; + + params.request_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.request_indices_offset); + params.qo_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_tile_indices_offset); + params.kv_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_tile_indices_offset); + params.o_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.o_indptr_offset); + params.kv_chunk_size_ptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_chunk_size_ptr_offset); + if (plan_info.split_kv) { + params.merge_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.merge_indptr_offset); + tmp_v = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.v_offset); + tmp_s = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.s_offset); + if (plan_info.enable_cuda_graph) { + params.block_valid_mask = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.block_valid_mask_offset); + } + } + params.padded_batch_size = plan_info.padded_batch_size; + params.max_total_num_rows = plan_info.total_num_rows; + if (plan_info.enable_cuda_graph) { + params.total_num_rows = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.total_num_rows_offset); + } + + cudaError_t status = cudaSuccess; + + DISPATCH_CTA_TILE_Q(plan_info.cta_tile_q, CTA_TILE_Q, { + status = flashinfer::BatchPrefillWithRaggedKVCacheDispatched< + CTA_TILE_Q, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + /*use_fp16_qk_reduction=*/USE_FP16_QK_REDUCTION, MASK_MODE, AttentionVariant, + RaggedParams>(params, tmp_v, tmp_s, enable_pdl, stream); + }); + + TORCH_CHECK(status == cudaSuccess, "BatchPrefillWithRaggedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); +} + +void BatchPrefillWithPagedKVCacheRun( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, + at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + PrefillPlanInfo plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + QKVLayout kv_layout = static_cast(layout); + auto device = q.device(); + int64_t batch_size = paged_kv_indptr.size(0) - 1; + int64_t num_qo_heads = q.size(1); + int64_t num_kv_heads, page_size; + uint32_t head_dim_qk = q.size(2); + if (kv_layout == QKVLayout::kHND) { + num_kv_heads = paged_k_cache.size(1); + page_size = paged_k_cache.size(2); + } else { + page_size = paged_k_cache.size(1); + num_kv_heads = paged_k_cache.size(2); + } + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == q.size(0), lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == q.size(1), lse.size(1), q.size(1)); + } + + void* float_buffer_ptr = static_cast(float_workspace_buffer.data_ptr()); + void* int_buffer_ptr = static_cast(int_workspace_buffer.data_ptr()); + + const MaskMode mask_mode = static_cast(mask_mode_code); + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = paged_k_cache.scalar_type(); + + // get q_stride_n and q_stride_h + const auto q_stride_n = q.stride(0); + const auto q_stride_h = q.stride(1); + + // get kv_cache_strides + const int64_t* kv_cache_strides = nullptr; + auto k_strides = paged_k_cache.strides(); + auto v_strides = paged_v_cache.strides(); + TORCH_CHECK(k_strides == v_strides, "k/v strides must be identical"); + kv_cache_strides = k_strides.data(); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, + RaggedParams, PagedParams, [&] { + PagedParams params; + + params.q = static_cast(q.data_ptr()); + paged_kv_t paged_kv( + num_kv_heads, page_size, HEAD_DIM_VO, batch_size, kv_layout, + static_cast(paged_k_cache.data_ptr()), + static_cast(paged_v_cache.data_ptr()), kv_cache_strides, + static_cast(paged_kv_indices.data_ptr()), + static_cast(paged_kv_indptr.data_ptr()), + static_cast(paged_kv_last_page_len.data_ptr())); + params.paged_kv = paged_kv; + params.q_indptr = static_cast(qo_indptr.data_ptr()); + params.o = static_cast(o.data_ptr()); + + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.num_qo_heads = num_qo_heads; + params.group_size = uint_fastdiv(num_qo_heads / paged_kv.num_heads); + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.window_left = window_left; + + params.request_indices = nullptr; + params.qo_tile_indices = nullptr; + params.kv_tile_indices = nullptr; + params.merge_indptr = nullptr; + params.o_indptr = nullptr; + params.kv_chunk_size_ptr = nullptr; + params.block_valid_mask = nullptr; + params.total_num_rows = nullptr; + params.max_total_num_rows = 0; + params.padded_batch_size = 0; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + DTypeO* tmp_v = nullptr; + float* tmp_s = nullptr; + + params.request_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.request_indices_offset); + params.qo_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_tile_indices_offset); + params.kv_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_tile_indices_offset); + params.o_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.o_indptr_offset); + params.kv_chunk_size_ptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_chunk_size_ptr_offset); + if (plan_info.split_kv) { + params.merge_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.merge_indptr_offset); + tmp_v = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.v_offset); + tmp_s = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.s_offset); + if (plan_info.enable_cuda_graph) { + params.block_valid_mask = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.block_valid_mask_offset); + } + } + params.padded_batch_size = plan_info.padded_batch_size; + params.max_total_num_rows = plan_info.total_num_rows; + if (plan_info.enable_cuda_graph) { + params.total_num_rows = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.total_num_rows_offset); + } + + cudaError_t status = cudaSuccess; + + DISPATCH_CTA_TILE_Q(plan_info.cta_tile_q, CTA_TILE_Q, { + status = flashinfer::BatchPrefillWithPagedKVCacheDispatched< + CTA_TILE_Q, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + /*use_fp16_qk_reduction=*/USE_FP16_QK_REDUCTION, MASK_MODE, AttentionVariant, + PagedParams>(params, tmp_v, tmp_s, enable_pdl, stream); + }); + + TORCH_CHECK(status == cudaSuccess, "BatchPrefillWithPagedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); +} diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_config.inc b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..3d563c083e2b971b91349b413d82d47e18fefcfc --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_config.inc @@ -0,0 +1,153 @@ +#pragma once +#include +#include +#include +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_custom_mask, std::optional maybe_mask_indptr, std::optional maybe_alibi_slopes, std::optional maybe_prefix_len_ptr, std::optional maybe_token_pos_in_items_ptr, std::optional maybe_max_item_len_ptr, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta, int64_t token_pos_in_items_len +#define ADDITIONAL_PARAMS_SETTER params.maybe_custom_mask = maybe_custom_mask ? static_cast(maybe_custom_mask->data_ptr()): nullptr; \ +params.maybe_mask_indptr = maybe_mask_indptr ? static_cast(maybe_mask_indptr->data_ptr()): nullptr; \ +params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.maybe_prefix_len_ptr = maybe_prefix_len_ptr ? static_cast(maybe_prefix_len_ptr->data_ptr()): nullptr; \ +params.maybe_token_pos_in_items_ptr = maybe_token_pos_in_items_ptr ? static_cast(maybe_token_pos_in_items_ptr->data_ptr()): nullptr; \ +params.maybe_max_item_len_ptr = maybe_max_item_len_ptr ? static_cast(maybe_max_item_len_ptr->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; \ +params.token_pos_in_items_len = token_pos_in_items_len; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, RaggedParams, PagedParams, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { \ + constexpr auto use_custom_mask = MASK_MODE == MaskMode::kCustom; \ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ + }) + +using namespace flashinfer; + +using DTypeQ = half; +using DTypeKV = half; +using DTypeO = half; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 64; +constexpr int HEAD_DIM_VO = 64; +constexpr bool USE_FP16_QK_REDUCTION = false; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = false; + + +struct RaggedParams { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + DTypeQ* q; + DTypeKV* k; + DTypeKV* v; + IdType* q_indptr; + IdType* kv_indptr; + DTypeO* o; + float* lse; + uint_fastdiv group_size; + + uint8_t* maybe_custom_mask; +int32_t* maybe_mask_indptr; +float* maybe_alibi_slopes; +uint32_t* maybe_prefix_len_ptr; +uint16_t* maybe_token_pos_in_items_ptr; +uint16_t* maybe_max_item_len_ptr; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; +int64_t token_pos_in_items_len; + + uint32_t num_qo_heads; + uint32_t num_kv_heads; + uint32_t q_stride_n; + uint32_t q_stride_h; + uint32_t k_stride_n; + uint32_t k_stride_h; + uint32_t v_stride_n; + uint32_t v_stride_h; + int32_t window_left; + + IdType* request_indices; + IdType* qo_tile_indices; + IdType* kv_tile_indices; + IdType* merge_indptr; + IdType* o_indptr; + IdType* kv_chunk_size_ptr; + bool* block_valid_mask; + uint32_t max_total_num_rows; + uint32_t* total_num_rows; + uint32_t padded_batch_size; + bool partition_kv; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { + return q_indptr[batch_idx + 1] - q_indptr[batch_idx]; + } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return kv_indptr[batch_idx + 1] - kv_indptr[batch_idx]; + } +}; + +struct PagedParams { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + DTypeQ* q; + paged_kv_t paged_kv; + IdType* q_indptr; + DTypeO* o; + float* lse; + uint_fastdiv group_size; + + uint8_t* maybe_custom_mask; +int32_t* maybe_mask_indptr; +float* maybe_alibi_slopes; +uint32_t* maybe_prefix_len_ptr; +uint16_t* maybe_token_pos_in_items_ptr; +uint16_t* maybe_max_item_len_ptr; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; +int64_t token_pos_in_items_len; + + uint32_t num_qo_heads; + IdType q_stride_n; + IdType q_stride_h; + int32_t window_left; + + IdType* request_indices; + IdType* qo_tile_indices; + IdType* kv_tile_indices; + IdType* merge_indptr; + IdType* o_indptr; + bool* block_valid_mask; + IdType* kv_chunk_size_ptr; + uint32_t max_total_num_rows; + uint32_t* total_num_rows; + uint32_t padded_batch_size; + bool partition_kv; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { + return q_indptr[batch_idx + 1] - q_indptr[batch_idx]; + } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return paged_kv.get_length(batch_idx); + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_jit_pybind.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..5421ab1cf910fdfbc867c182c8e22f23e44749ae --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_jit_pybind.cu @@ -0,0 +1,48 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "batch_prefill_config.inc" +#include "pytorch_extension_utils.h" + +at::Tensor BatchPrefillWithKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, + int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, + int64_t head_dim_vo, bool causal); + +void BatchPrefillWithRaggedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor k, at::Tensor v, + at::Tensor qo_indptr, at::Tensor kv_indptr, at::Tensor o, + std::optional maybe_lse, int64_t mask_mode_code, + int64_t layout, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +void BatchPrefillWithPagedKVCacheRun( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, + at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Batch-request prefill attention with KV-Cache plan + m.def("plan", BatchPrefillWithKVCachePlan); + // Batch-request prefill attention with KV-Cache operator + m.def("ragged_run", BatchPrefillWithRaggedKVCacheRun); + // Batch-request prefill attention with KV-Cache operator + m.def("paged_run", BatchPrefillWithPagedKVCacheRun); +} diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_0.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..0d1186c56489ce15e29d5f1f07d672cd68998352 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_0.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kNone == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_1.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..b155d2ddc8021119a99d691c9a5f60762a7389ea --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_1.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCausal == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_2.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..344bc67ac9289dcf2830826f8e82b2bfac986d3d --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_2.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCustom == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_3.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..ae7ea93585a9b50effde7e743ac9fe2322864644 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_3.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kMultiItemScoring == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_0.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..63093bf106f131fe01fc68450a946b34192bdbbe --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_0.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kNone == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_1.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..dc4ad3d120d3eefac6d2c821b1bf0462549b455f --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_1.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCausal == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_2.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..7c289d9f839d07256ffb3a6cdff27774130e28c9 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_2.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCustom == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_3.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..8df8e96108f270774806db09caefab6836b5446b --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_3.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kMultiItemScoring == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill.cu new file mode 100644 index 0000000000000000000000000000000000000000..a51fc7f56a0af7482bd5d2a9d3f672854df69dfd --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill.cu @@ -0,0 +1,330 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include +#include + +#include "batch_prefill_config.inc" +#include "pytorch_conversion_utils.h" +#include "pytorch_extension_utils.h" + +namespace flashinfer { + +template +cudaError_t BatchPrefillWithPagedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v, + float* tmp_s, bool enable_pdl, + cudaStream_t stream); + +template +cudaError_t BatchPrefillWithRaggedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v, + float* tmp_s, bool enable_pdl, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +at::Tensor BatchPrefillWithKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, + int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, + int64_t head_dim_vo, bool causal) { + size_t float_workspace_size_in_bytes = + float_workspace_buffer.size(0) * float_workspace_buffer.element_size(); + size_t int_workspace_size_in_bytes = + int_workspace_buffer.size(0) * int_workspace_buffer.element_size(); + + PrefillPlanInfo plan_info; + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + cudaError_t status = PrefillPlan( + float_workspace_buffer.data_ptr(), float_workspace_size_in_bytes, + int_workspace_buffer.data_ptr(), page_locked_int_workspace_buffer.data_ptr(), + int_workspace_size_in_bytes, plan_info, qo_indptr.data_ptr(), + kv_indptr.data_ptr(), total_num_rows, batch_size, num_qo_heads, num_kv_heads, + head_dim_qk, head_dim_vo, page_size, enable_cuda_graph, /*sizeof_dtype_o=*/2, stream); + + TORCH_CHECK(status == cudaSuccess, + "Failed to plan prefill with error: ", cudaGetErrorString(status)); + + return vec_to_tensor(plan_info.ToVector()); +} + +void BatchPrefillWithRaggedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor k, at::Tensor v, + at::Tensor qo_indptr, at::Tensor kv_indptr, at::Tensor o, + std::optional maybe_lse, int64_t mask_mode_code, + int64_t layout, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + PrefillPlanInfo plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + QKVLayout kv_layout = static_cast(layout); + + int64_t num_qo_heads = q.size(1); + int64_t head_dim_qk = q.size(2); + int64_t num_kv_heads = (kv_layout == QKVLayout::kNHD) ? k.size(1) : k.size(0); + uint32_t q_stride_n = q.stride(0), q_stride_h = q.stride(1), k_stride_n, k_stride_h, v_stride_n, + v_stride_h; + if (kv_layout == QKVLayout::kNHD) { + k_stride_n = k.stride(0); + k_stride_h = k.stride(1); + v_stride_n = v.stride(0); + v_stride_h = v.stride(1); + } else { + k_stride_h = k.stride(0); + k_stride_n = k.stride(1); + v_stride_h = v.stride(0); + v_stride_n = v.stride(1); + } + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == q.size(0), lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == q.size(1), lse.size(1), q.size(1)); + } + + void* float_buffer_ptr = float_workspace_buffer.data_ptr(); + void* int_buffer_ptr = int_workspace_buffer.data_ptr(); + + const MaskMode mask_mode = static_cast(mask_mode_code); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, + RaggedParams, PagedParams, [&] { + RaggedParams params; + + params.q = static_cast(q.data_ptr()); + params.k = static_cast(k.data_ptr()); + params.v = static_cast(v.data_ptr()); + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.q_indptr = static_cast(qo_indptr.data_ptr()); + params.kv_indptr = static_cast(kv_indptr.data_ptr()); + params.num_qo_heads = num_qo_heads; + params.num_kv_heads = num_kv_heads; + params.group_size = uint_fastdiv(num_qo_heads / num_kv_heads); + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.k_stride_n = k_stride_n; + params.k_stride_h = k_stride_h; + params.v_stride_n = v_stride_n; + params.v_stride_h = v_stride_h; + params.window_left = window_left; + + params.request_indices = nullptr; + params.qo_tile_indices = nullptr; + params.kv_tile_indices = nullptr; + params.merge_indptr = nullptr; + params.o_indptr = nullptr; + params.kv_chunk_size_ptr = nullptr; + params.block_valid_mask = nullptr; + params.total_num_rows = nullptr; + params.max_total_num_rows = 0; + params.padded_batch_size = 0; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + DTypeO* tmp_v = nullptr; + float* tmp_s = nullptr; + + params.request_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.request_indices_offset); + params.qo_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_tile_indices_offset); + params.kv_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_tile_indices_offset); + params.o_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.o_indptr_offset); + params.kv_chunk_size_ptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_chunk_size_ptr_offset); + if (plan_info.split_kv) { + params.merge_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.merge_indptr_offset); + tmp_v = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.v_offset); + tmp_s = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.s_offset); + if (plan_info.enable_cuda_graph) { + params.block_valid_mask = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.block_valid_mask_offset); + } + } + params.padded_batch_size = plan_info.padded_batch_size; + params.max_total_num_rows = plan_info.total_num_rows; + if (plan_info.enable_cuda_graph) { + params.total_num_rows = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.total_num_rows_offset); + } + + cudaError_t status = cudaSuccess; + + DISPATCH_CTA_TILE_Q(plan_info.cta_tile_q, CTA_TILE_Q, { + status = flashinfer::BatchPrefillWithRaggedKVCacheDispatched< + CTA_TILE_Q, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + /*use_fp16_qk_reduction=*/USE_FP16_QK_REDUCTION, MASK_MODE, AttentionVariant, + RaggedParams>(params, tmp_v, tmp_s, enable_pdl, stream); + }); + + TORCH_CHECK(status == cudaSuccess, "BatchPrefillWithRaggedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); +} + +void BatchPrefillWithPagedKVCacheRun( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, + at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS) { + PrefillPlanInfo plan_info; + plan_info.FromVector(tensor_to_vec(plan_info_vec)); + QKVLayout kv_layout = static_cast(layout); + auto device = q.device(); + int64_t batch_size = paged_kv_indptr.size(0) - 1; + int64_t num_qo_heads = q.size(1); + int64_t num_kv_heads, page_size; + uint32_t head_dim_qk = q.size(2); + if (kv_layout == QKVLayout::kHND) { + num_kv_heads = paged_k_cache.size(1); + page_size = paged_k_cache.size(2); + } else { + page_size = paged_k_cache.size(1); + num_kv_heads = paged_k_cache.size(2); + } + + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == q.size(0), lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == q.size(1), lse.size(1), q.size(1)); + } + + void* float_buffer_ptr = static_cast(float_workspace_buffer.data_ptr()); + void* int_buffer_ptr = static_cast(int_workspace_buffer.data_ptr()); + + const MaskMode mask_mode = static_cast(mask_mode_code); + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = paged_k_cache.scalar_type(); + + // get q_stride_n and q_stride_h + const auto q_stride_n = q.stride(0); + const auto q_stride_h = q.stride(1); + + // get kv_cache_strides + const int64_t* kv_cache_strides = nullptr; + auto k_strides = paged_k_cache.strides(); + auto v_strides = paged_v_cache.strides(); + TORCH_CHECK(k_strides == v_strides, "k/v strides must be identical"); + kv_cache_strides = k_strides.data(); + + const c10::cuda::OptionalCUDAGuard device_guard(float_workspace_buffer.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, + RaggedParams, PagedParams, [&] { + PagedParams params; + + params.q = static_cast(q.data_ptr()); + paged_kv_t paged_kv( + num_kv_heads, page_size, HEAD_DIM_VO, batch_size, kv_layout, + static_cast(paged_k_cache.data_ptr()), + static_cast(paged_v_cache.data_ptr()), kv_cache_strides, + static_cast(paged_kv_indices.data_ptr()), + static_cast(paged_kv_indptr.data_ptr()), + static_cast(paged_kv_last_page_len.data_ptr())); + params.paged_kv = paged_kv; + params.q_indptr = static_cast(qo_indptr.data_ptr()); + params.o = static_cast(o.data_ptr()); + + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.num_qo_heads = num_qo_heads; + params.group_size = uint_fastdiv(num_qo_heads / paged_kv.num_heads); + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.window_left = window_left; + + params.request_indices = nullptr; + params.qo_tile_indices = nullptr; + params.kv_tile_indices = nullptr; + params.merge_indptr = nullptr; + params.o_indptr = nullptr; + params.kv_chunk_size_ptr = nullptr; + params.block_valid_mask = nullptr; + params.total_num_rows = nullptr; + params.max_total_num_rows = 0; + params.padded_batch_size = 0; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + DTypeO* tmp_v = nullptr; + float* tmp_s = nullptr; + + params.request_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.request_indices_offset); + params.qo_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.qo_tile_indices_offset); + params.kv_tile_indices = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_tile_indices_offset); + params.o_indptr = GetPtrFromBaseOffset(int_buffer_ptr, plan_info.o_indptr_offset); + params.kv_chunk_size_ptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.kv_chunk_size_ptr_offset); + if (plan_info.split_kv) { + params.merge_indptr = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.merge_indptr_offset); + tmp_v = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.v_offset); + tmp_s = GetPtrFromBaseOffset(float_buffer_ptr, plan_info.s_offset); + if (plan_info.enable_cuda_graph) { + params.block_valid_mask = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.block_valid_mask_offset); + } + } + params.padded_batch_size = plan_info.padded_batch_size; + params.max_total_num_rows = plan_info.total_num_rows; + if (plan_info.enable_cuda_graph) { + params.total_num_rows = + GetPtrFromBaseOffset(int_buffer_ptr, plan_info.total_num_rows_offset); + } + + cudaError_t status = cudaSuccess; + + DISPATCH_CTA_TILE_Q(plan_info.cta_tile_q, CTA_TILE_Q, { + status = flashinfer::BatchPrefillWithPagedKVCacheDispatched< + CTA_TILE_Q, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + /*use_fp16_qk_reduction=*/USE_FP16_QK_REDUCTION, MASK_MODE, AttentionVariant, + PagedParams>(params, tmp_v, tmp_s, enable_pdl, stream); + }); + + TORCH_CHECK(status == cudaSuccess, "BatchPrefillWithPagedKVCache failed with error ", + cudaGetErrorString(status)); + return true; + }); +} diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_config.inc b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..5e60993f792075e65d4f29aa28a492a0a45ddfff --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_config.inc @@ -0,0 +1,153 @@ +#pragma once +#include +#include +#include +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_custom_mask, std::optional maybe_mask_indptr, std::optional maybe_alibi_slopes, std::optional maybe_prefix_len_ptr, std::optional maybe_token_pos_in_items_ptr, std::optional maybe_max_item_len_ptr, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta, int64_t token_pos_in_items_len +#define ADDITIONAL_PARAMS_SETTER params.maybe_custom_mask = maybe_custom_mask ? static_cast(maybe_custom_mask->data_ptr()): nullptr; \ +params.maybe_mask_indptr = maybe_mask_indptr ? static_cast(maybe_mask_indptr->data_ptr()): nullptr; \ +params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.maybe_prefix_len_ptr = maybe_prefix_len_ptr ? static_cast(maybe_prefix_len_ptr->data_ptr()): nullptr; \ +params.maybe_token_pos_in_items_ptr = maybe_token_pos_in_items_ptr ? static_cast(maybe_token_pos_in_items_ptr->data_ptr()): nullptr; \ +params.maybe_max_item_len_ptr = maybe_max_item_len_ptr ? static_cast(maybe_max_item_len_ptr->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; \ +params.token_pos_in_items_len = token_pos_in_items_len; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, RaggedParams, PagedParams, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { \ + constexpr auto use_custom_mask = MASK_MODE == MaskMode::kCustom; \ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ + }) + +using namespace flashinfer; + +using DTypeQ = half; +using DTypeKV = half; +using DTypeO = half; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 64; +constexpr int HEAD_DIM_VO = 64; +constexpr bool USE_FP16_QK_REDUCTION = false; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = true; + + +struct RaggedParams { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + DTypeQ* q; + DTypeKV* k; + DTypeKV* v; + IdType* q_indptr; + IdType* kv_indptr; + DTypeO* o; + float* lse; + uint_fastdiv group_size; + + uint8_t* maybe_custom_mask; +int32_t* maybe_mask_indptr; +float* maybe_alibi_slopes; +uint32_t* maybe_prefix_len_ptr; +uint16_t* maybe_token_pos_in_items_ptr; +uint16_t* maybe_max_item_len_ptr; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; +int64_t token_pos_in_items_len; + + uint32_t num_qo_heads; + uint32_t num_kv_heads; + uint32_t q_stride_n; + uint32_t q_stride_h; + uint32_t k_stride_n; + uint32_t k_stride_h; + uint32_t v_stride_n; + uint32_t v_stride_h; + int32_t window_left; + + IdType* request_indices; + IdType* qo_tile_indices; + IdType* kv_tile_indices; + IdType* merge_indptr; + IdType* o_indptr; + IdType* kv_chunk_size_ptr; + bool* block_valid_mask; + uint32_t max_total_num_rows; + uint32_t* total_num_rows; + uint32_t padded_batch_size; + bool partition_kv; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { + return q_indptr[batch_idx + 1] - q_indptr[batch_idx]; + } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return kv_indptr[batch_idx + 1] - kv_indptr[batch_idx]; + } +}; + +struct PagedParams { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + DTypeQ* q; + paged_kv_t paged_kv; + IdType* q_indptr; + DTypeO* o; + float* lse; + uint_fastdiv group_size; + + uint8_t* maybe_custom_mask; +int32_t* maybe_mask_indptr; +float* maybe_alibi_slopes; +uint32_t* maybe_prefix_len_ptr; +uint16_t* maybe_token_pos_in_items_ptr; +uint16_t* maybe_max_item_len_ptr; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; +int64_t token_pos_in_items_len; + + uint32_t num_qo_heads; + IdType q_stride_n; + IdType q_stride_h; + int32_t window_left; + + IdType* request_indices; + IdType* qo_tile_indices; + IdType* kv_tile_indices; + IdType* merge_indptr; + IdType* o_indptr; + bool* block_valid_mask; + IdType* kv_chunk_size_ptr; + uint32_t max_total_num_rows; + uint32_t* total_num_rows; + uint32_t padded_batch_size; + bool partition_kv; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { + return q_indptr[batch_idx + 1] - q_indptr[batch_idx]; + } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return paged_kv.get_length(batch_idx); + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_jit_pybind.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..5421ab1cf910fdfbc867c182c8e22f23e44749ae --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_jit_pybind.cu @@ -0,0 +1,48 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "batch_prefill_config.inc" +#include "pytorch_extension_utils.h" + +at::Tensor BatchPrefillWithKVCachePlan( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, + at::Tensor page_locked_int_workspace_buffer, at::Tensor qo_indptr, at::Tensor kv_indptr, + at::Tensor kv_len_arr, int64_t total_num_rows, int64_t batch_size, int64_t num_qo_heads, + int64_t num_kv_heads, int64_t page_size, bool enable_cuda_graph, int64_t head_dim_qk, + int64_t head_dim_vo, bool causal); + +void BatchPrefillWithRaggedKVCacheRun(at::Tensor float_workspace_buffer, + at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor k, at::Tensor v, + at::Tensor qo_indptr, at::Tensor kv_indptr, at::Tensor o, + std::optional maybe_lse, int64_t mask_mode_code, + int64_t layout, int64_t window_left, + bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +void BatchPrefillWithPagedKVCacheRun( + at::Tensor float_workspace_buffer, at::Tensor int_workspace_buffer, at::Tensor plan_info_vec, + at::Tensor q, at::Tensor paged_k_cache, at::Tensor paged_v_cache, at::Tensor qo_indptr, + at::Tensor paged_kv_indptr, at::Tensor paged_kv_indices, at::Tensor paged_kv_last_page_len, + at::Tensor o, std::optional maybe_lse, int64_t mask_mode_code, int64_t layout, + int64_t window_left, bool enable_pdl ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Batch-request prefill attention with KV-Cache plan + m.def("plan", BatchPrefillWithKVCachePlan); + // Batch-request prefill attention with KV-Cache operator + m.def("ragged_run", BatchPrefillWithRaggedKVCacheRun); + // Batch-request prefill attention with KV-Cache operator + m.def("paged_run", BatchPrefillWithPagedKVCacheRun); +} diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_0.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..54b583ba709acb7db72c4e1854f934ef3fc54872 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_0.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kNone == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_1.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..ffceb828ae6798d6d4cd55042d9578f9924157b3 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_1.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCausal == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_2.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..1e52ae59b60c66d2eb32d4ec8c1eb829c9153df0 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_2.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCustom == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_3.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..170d1a6e67e5d2791d29cdc46d28e6755a237e44 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_paged_kernel_mask_3.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kMultiItemScoring == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithPagedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, PagedParams>(PagedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_0.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..96536df05b7b979cf5f79585eb7c4c9ae5ba2bbc --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_0.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kNone == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_1.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..8c513ea2d6fdec8108b49c8e94da978831075ddb --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_1.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCausal == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_2.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..b2e25b8da455088c2e73d4002cbf78824c0dfc39 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_2.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCustom == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_3.cu b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..6dd8e0b2a9973c9b006547e05016f1eba273a485 --- /dev/null +++ b/csrc/generated/batch_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_dtype_idx_i32_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/batch_prefill_ragged_kernel_mask_3.cu @@ -0,0 +1,22 @@ +#include +#include "batch_prefill_config.inc" + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kMultiItemScoring == MaskMode::kCustom; + + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/16, 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/64, 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + +template cudaError_t BatchPrefillWithRaggedKVCacheDispatched< + /*CTA_TILE_Q=*/128, 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, + DefaultAttention, RaggedParams>(RaggedParams params, half* tmp_v, float* tmp_s, bool enable_pdl, cudaStream_t stream); + + +}; // namespace flashinfer \ No newline at end of file diff --git a/csrc/generated/gelu_and_mul.cu b/csrc/generated/gelu_and_mul.cu new file mode 100644 index 0000000000000000000000000000000000000000..6f2f67ba2d0fd40f8d110593bb8bd556aede117f --- /dev/null +++ b/csrc/generated/gelu_and_mul.cu @@ -0,0 +1,51 @@ + +#include +#include "pytorch_extension_utils.h" +#include + + + +using namespace flashinfer; + + +__device__ __forceinline__ float gelu(const float& val) { + constexpr float kAlpha = M_SQRT1_2; + return val * 0.5f * (1.0f + ::erf(val * kAlpha)); +} + + +void gelu_and_mul(at::Tensor& out, at::Tensor& input, bool enable_pdl) { + int d = input.size(-1) / 2; + int64_t num_tokens = input.numel() / input.size(-1); + dim3 grid(num_tokens); + + const c10::cuda::OptionalCUDAGuard device_guard(out.device()); + auto stream = at::cuda::getCurrentCUDAStream(); + DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FP16(input.scalar_type(), c_type, [&] { + uint32_t vec_size = 16 / sizeof(c_type); + cudaLaunchConfig_t config; + config.gridDim = num_tokens; + config.blockDim = std::min(d / vec_size, 1024U); + config.dynamicSmemBytes = 0; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = enable_pdl; + config.numAttrs = 1; + config.attrs = attrs; + + auto kernel = flashinfer::activation::act_and_mul_kernel; + + cudaLaunchKernelEx(&config, kernel, static_cast(out.data_ptr()), + static_cast(input.data_ptr()), d); + + cudaError_t err = cudaGetLastError(); + TORCH_CHECK(err == cudaSuccess, "Failed to launch kernel: ", cudaGetErrorString(err)); + + return true; + }); +} + +// TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { +// m.def("gelu_and_mul", gelu_and_mul); +// } \ No newline at end of file diff --git a/csrc/generated/gelu_tanh_and_mul.cu b/csrc/generated/gelu_tanh_and_mul.cu new file mode 100644 index 0000000000000000000000000000000000000000..13139cb7bfb5f061ca2f21fd189f00e975096a53 --- /dev/null +++ b/csrc/generated/gelu_tanh_and_mul.cu @@ -0,0 +1,52 @@ + +#include +#include "pytorch_extension_utils.h" +#include + + + +using namespace flashinfer; + + +__device__ __forceinline__ float gelu_tanh(const float& val) { + const float cdf = + 0.5f * (1.0f + math::tanh((0.7978845608028654f * (val + 0.044715f * val * val * val)))); + return val * cdf; +} + + +void gelu_tanh_and_mul(at::Tensor& out, at::Tensor& input, bool enable_pdl) { + int d = input.size(-1) / 2; + int64_t num_tokens = input.numel() / input.size(-1); + dim3 grid(num_tokens); + + const c10::cuda::OptionalCUDAGuard device_guard(out.device()); + auto stream = at::cuda::getCurrentCUDAStream(); + DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FP16(input.scalar_type(), c_type, [&] { + uint32_t vec_size = 16 / sizeof(c_type); + cudaLaunchConfig_t config; + config.gridDim = num_tokens; + config.blockDim = std::min(d / vec_size, 1024U); + config.dynamicSmemBytes = 0; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = enable_pdl; + config.numAttrs = 1; + config.attrs = attrs; + + auto kernel = flashinfer::activation::act_and_mul_kernel; + + cudaLaunchKernelEx(&config, kernel, static_cast(out.data_ptr()), + static_cast(input.data_ptr()), d); + + cudaError_t err = cudaGetLastError(); + TORCH_CHECK(err == cudaSuccess, "Failed to launch kernel: ", cudaGetErrorString(err)); + + return true; + }); +} + +// TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { +// m.def("gelu_tanh_and_mul", gelu_tanh_and_mul); +// } \ No newline at end of file diff --git a/csrc/generated/silu_and_mul.cu b/csrc/generated/silu_and_mul.cu new file mode 100644 index 0000000000000000000000000000000000000000..b342b9cf473ac9a97f61f3de2dd79cd29b6bcde7 --- /dev/null +++ b/csrc/generated/silu_and_mul.cu @@ -0,0 +1,50 @@ + +#include +#include "pytorch_extension_utils.h" +#include + + + +using namespace flashinfer; + + +__device__ __forceinline__ float silu(const float& val) { + return val / (1.0f + __expf(-val)); +} + + +void silu_and_mul(at::Tensor& out, at::Tensor& input, bool enable_pdl) { + int d = input.size(-1) / 2; + int64_t num_tokens = input.numel() / input.size(-1); + dim3 grid(num_tokens); + + const c10::cuda::OptionalCUDAGuard device_guard(out.device()); + auto stream = at::cuda::getCurrentCUDAStream(); + DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FP16(input.scalar_type(), c_type, [&] { + uint32_t vec_size = 16 / sizeof(c_type); + cudaLaunchConfig_t config; + config.gridDim = num_tokens; + config.blockDim = std::min(d / vec_size, 1024U); + config.dynamicSmemBytes = 0; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = enable_pdl; + config.numAttrs = 1; + config.attrs = attrs; + + auto kernel = flashinfer::activation::act_and_mul_kernel; + + cudaLaunchKernelEx(&config, kernel, static_cast(out.data_ptr()), + static_cast(input.data_ptr()), d); + + cudaError_t err = cudaGetLastError(); + TORCH_CHECK(err == cudaSuccess, "Failed to launch kernel: ", cudaGetErrorString(err)); + + return true; + }); +} + +// TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { +// m.def("silu_and_mul", silu_and_mul); +// } \ No newline at end of file diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/single_decode.cu b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/single_decode.cu new file mode 100644 index 0000000000000000000000000000000000000000..a22ebd91727745cd8a85beea286d3c80bef3128c --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/single_decode.cu @@ -0,0 +1,104 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include + +#include "pytorch_extension_utils.h" +#include "single_decode_config.inc" + +namespace flashinfer { + +template +cudaError_t SingleDecodeWithKVCacheDispatched(Params params, typename Params::DTypeO* tmp, + cudaStream_t stream); +} // namespace flashinfer + +using namespace flashinfer; + +void single_decode_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS) { + CHECK_INPUT(q); + CHECK_INPUT(k); + CHECK_INPUT(v); + CHECK_INPUT(tmp); + auto device = q.device(); + CHECK_EQ(k.device(), device); + CHECK_EQ(v.device(), device); + CHECK_EQ(tmp.device(), device); + CHECK_DIM(2, q); + CHECK_DIM(3, k); + CHECK_DIM(3, v); + CHECK_SHAPE(k, v); + CHECK_EQ(q.size(1), k.size(2)); + CHECK_EQ(v.scalar_type(), k.scalar_type()); + unsigned int num_qo_heads = q.size(0); + unsigned int head_dim_qk = q.size(1); + unsigned int head_dim_vo = v.size(2); + unsigned int kv_len, num_kv_heads; + QKVLayout kv_layout = static_cast(layout); + if (kv_layout == QKVLayout::kNHD) { + kv_len = k.size(0); + num_kv_heads = k.size(1); + } else { + num_kv_heads = k.size(0); + kv_len = k.size(1); + } + CHECK_GQA_HEAD_DIVISIBLE(num_qo_heads, num_kv_heads); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(device); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + TORCH_CHECK(head_dim_qk == head_dim_vo, + "CUDA cores template only supports equal head dim for QK and VO, please use tensor " + "cores template for different head dim"); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + Params params; + + params.q = static_cast(q.data_ptr()); + params.k = static_cast(k.data_ptr()); + params.v = static_cast(v.data_ptr()); + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.kv_len = kv_len; + params.num_qo_heads = num_qo_heads; + params.num_kv_heads = num_kv_heads; + params.q_stride_n = num_qo_heads * head_dim_qk; + params.q_stride_h = head_dim_qk; + params.kv_stride_n = + (kv_layout == QKVLayout::kNHD) ? num_kv_heads * head_dim_vo : head_dim_vo; + params.kv_stride_h = (kv_layout == QKVLayout::kNHD) ? head_dim_vo : kv_len * head_dim_vo; + params.window_left = window_left; + params.kv_chunk_size = 0; + + ADDITIONAL_PARAMS_SETTER + + cudaError_t status = + flashinfer::SingleDecodeWithKVCacheDispatched( + params, static_cast(tmp.data_ptr()), stream); + TORCH_CHECK(status == cudaSuccess, "SingleDecodeWithKVCache kernel launch failed, error: " + + std::string(cudaGetErrorString(status))); + return true; + }); +} diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/single_decode_config.inc b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/single_decode_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..d2ffe8ac286ea7cccfcdcf302116f1a00152ec43 --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/single_decode_config.inc @@ -0,0 +1,64 @@ +#pragma once +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta +#define ADDITIONAL_PARAMS_SETTER params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, ...) {\ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ +} + +using namespace flashinfer; + +using DTypeQ = nv_bfloat16; +using DTypeKV = nv_bfloat16; +using DTypeO = nv_bfloat16; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 128; +constexpr int HEAD_DIM_VO = 128; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = false; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = int32_t; + DTypeQ* q; + DTypeKV* k; + DTypeKV* v; + DTypeO* o; + float* lse; + float* maybe_alibi_slopes; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; + + uint32_t kv_len; + uint32_t num_qo_heads; + uint32_t num_kv_heads; + uint32_t q_stride_n; + uint32_t q_stride_h; + uint32_t kv_stride_n; + uint32_t kv_stride_h; + int32_t window_left; + uint32_t kv_chunk_size; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { return 1; } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return kv_len; + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/single_decode_jit_pybind.cu b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/single_decode_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..f4c360d3e4d283239d95810d4b5aa5bbd05cb953 --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/single_decode_jit_pybind.cu @@ -0,0 +1,27 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "pytorch_extension_utils.h" +#include "single_decode_config.inc" + +void single_decode_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Single-request decode with KV-Cache operator + m.def("run", single_decode_with_kv_cache); +} diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/single_decode_kernel.cu b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/single_decode_kernel.cu new file mode 100644 index 0000000000000000000000000000000000000000..c2ec47236fa827ed1ae4272d6a80a852d73ff5f4 --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/single_decode_kernel.cu @@ -0,0 +1,13 @@ +#include +#include "single_decode_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SingleDecodeWithKVCacheDispatched< + 128, PosEncodingMode::kNone, DefaultAttention, Params>( + Params params, nv_bfloat16* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/single_decode.cu b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/single_decode.cu new file mode 100644 index 0000000000000000000000000000000000000000..a22ebd91727745cd8a85beea286d3c80bef3128c --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/single_decode.cu @@ -0,0 +1,104 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include + +#include "pytorch_extension_utils.h" +#include "single_decode_config.inc" + +namespace flashinfer { + +template +cudaError_t SingleDecodeWithKVCacheDispatched(Params params, typename Params::DTypeO* tmp, + cudaStream_t stream); +} // namespace flashinfer + +using namespace flashinfer; + +void single_decode_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS) { + CHECK_INPUT(q); + CHECK_INPUT(k); + CHECK_INPUT(v); + CHECK_INPUT(tmp); + auto device = q.device(); + CHECK_EQ(k.device(), device); + CHECK_EQ(v.device(), device); + CHECK_EQ(tmp.device(), device); + CHECK_DIM(2, q); + CHECK_DIM(3, k); + CHECK_DIM(3, v); + CHECK_SHAPE(k, v); + CHECK_EQ(q.size(1), k.size(2)); + CHECK_EQ(v.scalar_type(), k.scalar_type()); + unsigned int num_qo_heads = q.size(0); + unsigned int head_dim_qk = q.size(1); + unsigned int head_dim_vo = v.size(2); + unsigned int kv_len, num_kv_heads; + QKVLayout kv_layout = static_cast(layout); + if (kv_layout == QKVLayout::kNHD) { + kv_len = k.size(0); + num_kv_heads = k.size(1); + } else { + num_kv_heads = k.size(0); + kv_len = k.size(1); + } + CHECK_GQA_HEAD_DIVISIBLE(num_qo_heads, num_kv_heads); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(device); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + TORCH_CHECK(head_dim_qk == head_dim_vo, + "CUDA cores template only supports equal head dim for QK and VO, please use tensor " + "cores template for different head dim"); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + Params params; + + params.q = static_cast(q.data_ptr()); + params.k = static_cast(k.data_ptr()); + params.v = static_cast(v.data_ptr()); + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.kv_len = kv_len; + params.num_qo_heads = num_qo_heads; + params.num_kv_heads = num_kv_heads; + params.q_stride_n = num_qo_heads * head_dim_qk; + params.q_stride_h = head_dim_qk; + params.kv_stride_n = + (kv_layout == QKVLayout::kNHD) ? num_kv_heads * head_dim_vo : head_dim_vo; + params.kv_stride_h = (kv_layout == QKVLayout::kNHD) ? head_dim_vo : kv_len * head_dim_vo; + params.window_left = window_left; + params.kv_chunk_size = 0; + + ADDITIONAL_PARAMS_SETTER + + cudaError_t status = + flashinfer::SingleDecodeWithKVCacheDispatched( + params, static_cast(tmp.data_ptr()), stream); + TORCH_CHECK(status == cudaSuccess, "SingleDecodeWithKVCache kernel launch failed, error: " + + std::string(cudaGetErrorString(status))); + return true; + }); +} diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/single_decode_config.inc b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/single_decode_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..c304da6ae405aec165e36cd8af93c1a1124b86bd --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/single_decode_config.inc @@ -0,0 +1,64 @@ +#pragma once +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta +#define ADDITIONAL_PARAMS_SETTER params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, ...) {\ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ +} + +using namespace flashinfer; + +using DTypeQ = nv_bfloat16; +using DTypeKV = nv_bfloat16; +using DTypeO = nv_bfloat16; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 256; +constexpr int HEAD_DIM_VO = 256; +constexpr auto USE_LOGITS_SOFT_CAP = true; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = true; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = int32_t; + DTypeQ* q; + DTypeKV* k; + DTypeKV* v; + DTypeO* o; + float* lse; + float* maybe_alibi_slopes; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; + + uint32_t kv_len; + uint32_t num_qo_heads; + uint32_t num_kv_heads; + uint32_t q_stride_n; + uint32_t q_stride_h; + uint32_t kv_stride_n; + uint32_t kv_stride_h; + int32_t window_left; + uint32_t kv_chunk_size; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { return 1; } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return kv_len; + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/single_decode_jit_pybind.cu b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/single_decode_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..f4c360d3e4d283239d95810d4b5aa5bbd05cb953 --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/single_decode_jit_pybind.cu @@ -0,0 +1,27 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "pytorch_extension_utils.h" +#include "single_decode_config.inc" + +void single_decode_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Single-request decode with KV-Cache operator + m.def("run", single_decode_with_kv_cache); +} diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/single_decode_kernel.cu b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/single_decode_kernel.cu new file mode 100644 index 0000000000000000000000000000000000000000..befd17566ba95c132a8a3fd19a2b756a59ada350 --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/single_decode_kernel.cu @@ -0,0 +1,13 @@ +#include +#include "single_decode_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SingleDecodeWithKVCacheDispatched< + 256, PosEncodingMode::kNone, DefaultAttention, Params>( + Params params, nv_bfloat16* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/single_decode.cu b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/single_decode.cu new file mode 100644 index 0000000000000000000000000000000000000000..a22ebd91727745cd8a85beea286d3c80bef3128c --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/single_decode.cu @@ -0,0 +1,104 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include + +#include "pytorch_extension_utils.h" +#include "single_decode_config.inc" + +namespace flashinfer { + +template +cudaError_t SingleDecodeWithKVCacheDispatched(Params params, typename Params::DTypeO* tmp, + cudaStream_t stream); +} // namespace flashinfer + +using namespace flashinfer; + +void single_decode_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS) { + CHECK_INPUT(q); + CHECK_INPUT(k); + CHECK_INPUT(v); + CHECK_INPUT(tmp); + auto device = q.device(); + CHECK_EQ(k.device(), device); + CHECK_EQ(v.device(), device); + CHECK_EQ(tmp.device(), device); + CHECK_DIM(2, q); + CHECK_DIM(3, k); + CHECK_DIM(3, v); + CHECK_SHAPE(k, v); + CHECK_EQ(q.size(1), k.size(2)); + CHECK_EQ(v.scalar_type(), k.scalar_type()); + unsigned int num_qo_heads = q.size(0); + unsigned int head_dim_qk = q.size(1); + unsigned int head_dim_vo = v.size(2); + unsigned int kv_len, num_kv_heads; + QKVLayout kv_layout = static_cast(layout); + if (kv_layout == QKVLayout::kNHD) { + kv_len = k.size(0); + num_kv_heads = k.size(1); + } else { + num_kv_heads = k.size(0); + kv_len = k.size(1); + } + CHECK_GQA_HEAD_DIVISIBLE(num_qo_heads, num_kv_heads); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(device); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + TORCH_CHECK(head_dim_qk == head_dim_vo, + "CUDA cores template only supports equal head dim for QK and VO, please use tensor " + "cores template for different head dim"); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + Params params; + + params.q = static_cast(q.data_ptr()); + params.k = static_cast(k.data_ptr()); + params.v = static_cast(v.data_ptr()); + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.kv_len = kv_len; + params.num_qo_heads = num_qo_heads; + params.num_kv_heads = num_kv_heads; + params.q_stride_n = num_qo_heads * head_dim_qk; + params.q_stride_h = head_dim_qk; + params.kv_stride_n = + (kv_layout == QKVLayout::kNHD) ? num_kv_heads * head_dim_vo : head_dim_vo; + params.kv_stride_h = (kv_layout == QKVLayout::kNHD) ? head_dim_vo : kv_len * head_dim_vo; + params.window_left = window_left; + params.kv_chunk_size = 0; + + ADDITIONAL_PARAMS_SETTER + + cudaError_t status = + flashinfer::SingleDecodeWithKVCacheDispatched( + params, static_cast(tmp.data_ptr()), stream); + TORCH_CHECK(status == cudaSuccess, "SingleDecodeWithKVCache kernel launch failed, error: " + + std::string(cudaGetErrorString(status))); + return true; + }); +} diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/single_decode_config.inc b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/single_decode_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..8d0b40b47772e33eeb70aab0c607cf0b104097a5 --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/single_decode_config.inc @@ -0,0 +1,64 @@ +#pragma once +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta +#define ADDITIONAL_PARAMS_SETTER params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, ...) {\ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ +} + +using namespace flashinfer; + +using DTypeQ = nv_bfloat16; +using DTypeKV = nv_bfloat16; +using DTypeO = nv_bfloat16; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 64; +constexpr int HEAD_DIM_VO = 64; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = false; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = int32_t; + DTypeQ* q; + DTypeKV* k; + DTypeKV* v; + DTypeO* o; + float* lse; + float* maybe_alibi_slopes; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; + + uint32_t kv_len; + uint32_t num_qo_heads; + uint32_t num_kv_heads; + uint32_t q_stride_n; + uint32_t q_stride_h; + uint32_t kv_stride_n; + uint32_t kv_stride_h; + int32_t window_left; + uint32_t kv_chunk_size; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { return 1; } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return kv_len; + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/single_decode_jit_pybind.cu b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/single_decode_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..f4c360d3e4d283239d95810d4b5aa5bbd05cb953 --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/single_decode_jit_pybind.cu @@ -0,0 +1,27 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "pytorch_extension_utils.h" +#include "single_decode_config.inc" + +void single_decode_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Single-request decode with KV-Cache operator + m.def("run", single_decode_with_kv_cache); +} diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/single_decode_kernel.cu b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/single_decode_kernel.cu new file mode 100644 index 0000000000000000000000000000000000000000..a72a34cca58a670438ef6d99276a121bc4cd0ef5 --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/single_decode_kernel.cu @@ -0,0 +1,13 @@ +#include +#include "single_decode_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SingleDecodeWithKVCacheDispatched< + 64, PosEncodingMode::kNone, DefaultAttention, Params>( + Params params, nv_bfloat16* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/single_decode.cu b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/single_decode.cu new file mode 100644 index 0000000000000000000000000000000000000000..a22ebd91727745cd8a85beea286d3c80bef3128c --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/single_decode.cu @@ -0,0 +1,104 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include + +#include "pytorch_extension_utils.h" +#include "single_decode_config.inc" + +namespace flashinfer { + +template +cudaError_t SingleDecodeWithKVCacheDispatched(Params params, typename Params::DTypeO* tmp, + cudaStream_t stream); +} // namespace flashinfer + +using namespace flashinfer; + +void single_decode_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS) { + CHECK_INPUT(q); + CHECK_INPUT(k); + CHECK_INPUT(v); + CHECK_INPUT(tmp); + auto device = q.device(); + CHECK_EQ(k.device(), device); + CHECK_EQ(v.device(), device); + CHECK_EQ(tmp.device(), device); + CHECK_DIM(2, q); + CHECK_DIM(3, k); + CHECK_DIM(3, v); + CHECK_SHAPE(k, v); + CHECK_EQ(q.size(1), k.size(2)); + CHECK_EQ(v.scalar_type(), k.scalar_type()); + unsigned int num_qo_heads = q.size(0); + unsigned int head_dim_qk = q.size(1); + unsigned int head_dim_vo = v.size(2); + unsigned int kv_len, num_kv_heads; + QKVLayout kv_layout = static_cast(layout); + if (kv_layout == QKVLayout::kNHD) { + kv_len = k.size(0); + num_kv_heads = k.size(1); + } else { + num_kv_heads = k.size(0); + kv_len = k.size(1); + } + CHECK_GQA_HEAD_DIVISIBLE(num_qo_heads, num_kv_heads); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(device); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + TORCH_CHECK(head_dim_qk == head_dim_vo, + "CUDA cores template only supports equal head dim for QK and VO, please use tensor " + "cores template for different head dim"); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + Params params; + + params.q = static_cast(q.data_ptr()); + params.k = static_cast(k.data_ptr()); + params.v = static_cast(v.data_ptr()); + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.kv_len = kv_len; + params.num_qo_heads = num_qo_heads; + params.num_kv_heads = num_kv_heads; + params.q_stride_n = num_qo_heads * head_dim_qk; + params.q_stride_h = head_dim_qk; + params.kv_stride_n = + (kv_layout == QKVLayout::kNHD) ? num_kv_heads * head_dim_vo : head_dim_vo; + params.kv_stride_h = (kv_layout == QKVLayout::kNHD) ? head_dim_vo : kv_len * head_dim_vo; + params.window_left = window_left; + params.kv_chunk_size = 0; + + ADDITIONAL_PARAMS_SETTER + + cudaError_t status = + flashinfer::SingleDecodeWithKVCacheDispatched( + params, static_cast(tmp.data_ptr()), stream); + TORCH_CHECK(status == cudaSuccess, "SingleDecodeWithKVCache kernel launch failed, error: " + + std::string(cudaGetErrorString(status))); + return true; + }); +} diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/single_decode_config.inc b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/single_decode_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..2f5ef8dd46f1dc4bfc9aa9cae67a6299f001fabb --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/single_decode_config.inc @@ -0,0 +1,64 @@ +#pragma once +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta +#define ADDITIONAL_PARAMS_SETTER params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, ...) {\ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ +} + +using namespace flashinfer; + +using DTypeQ = nv_bfloat16; +using DTypeKV = nv_bfloat16; +using DTypeO = nv_bfloat16; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 64; +constexpr int HEAD_DIM_VO = 64; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = true; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = int32_t; + DTypeQ* q; + DTypeKV* k; + DTypeKV* v; + DTypeO* o; + float* lse; + float* maybe_alibi_slopes; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; + + uint32_t kv_len; + uint32_t num_qo_heads; + uint32_t num_kv_heads; + uint32_t q_stride_n; + uint32_t q_stride_h; + uint32_t kv_stride_n; + uint32_t kv_stride_h; + int32_t window_left; + uint32_t kv_chunk_size; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { return 1; } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return kv_len; + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/single_decode_jit_pybind.cu b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/single_decode_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..f4c360d3e4d283239d95810d4b5aa5bbd05cb953 --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/single_decode_jit_pybind.cu @@ -0,0 +1,27 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "pytorch_extension_utils.h" +#include "single_decode_config.inc" + +void single_decode_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Single-request decode with KV-Cache operator + m.def("run", single_decode_with_kv_cache); +} diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/single_decode_kernel.cu b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/single_decode_kernel.cu new file mode 100644 index 0000000000000000000000000000000000000000..eff8599261be9b92499ce5fdadcd54596fac35cc --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/single_decode_kernel.cu @@ -0,0 +1,13 @@ +#include +#include "single_decode_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SingleDecodeWithKVCacheDispatched< + 64, PosEncodingMode::kNone, DefaultAttention, Params>( + Params params, nv_bfloat16* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/single_decode.cu b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/single_decode.cu new file mode 100644 index 0000000000000000000000000000000000000000..a22ebd91727745cd8a85beea286d3c80bef3128c --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/single_decode.cu @@ -0,0 +1,104 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include + +#include "pytorch_extension_utils.h" +#include "single_decode_config.inc" + +namespace flashinfer { + +template +cudaError_t SingleDecodeWithKVCacheDispatched(Params params, typename Params::DTypeO* tmp, + cudaStream_t stream); +} // namespace flashinfer + +using namespace flashinfer; + +void single_decode_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS) { + CHECK_INPUT(q); + CHECK_INPUT(k); + CHECK_INPUT(v); + CHECK_INPUT(tmp); + auto device = q.device(); + CHECK_EQ(k.device(), device); + CHECK_EQ(v.device(), device); + CHECK_EQ(tmp.device(), device); + CHECK_DIM(2, q); + CHECK_DIM(3, k); + CHECK_DIM(3, v); + CHECK_SHAPE(k, v); + CHECK_EQ(q.size(1), k.size(2)); + CHECK_EQ(v.scalar_type(), k.scalar_type()); + unsigned int num_qo_heads = q.size(0); + unsigned int head_dim_qk = q.size(1); + unsigned int head_dim_vo = v.size(2); + unsigned int kv_len, num_kv_heads; + QKVLayout kv_layout = static_cast(layout); + if (kv_layout == QKVLayout::kNHD) { + kv_len = k.size(0); + num_kv_heads = k.size(1); + } else { + num_kv_heads = k.size(0); + kv_len = k.size(1); + } + CHECK_GQA_HEAD_DIVISIBLE(num_qo_heads, num_kv_heads); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(device); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + TORCH_CHECK(head_dim_qk == head_dim_vo, + "CUDA cores template only supports equal head dim for QK and VO, please use tensor " + "cores template for different head dim"); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + Params params; + + params.q = static_cast(q.data_ptr()); + params.k = static_cast(k.data_ptr()); + params.v = static_cast(v.data_ptr()); + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.kv_len = kv_len; + params.num_qo_heads = num_qo_heads; + params.num_kv_heads = num_kv_heads; + params.q_stride_n = num_qo_heads * head_dim_qk; + params.q_stride_h = head_dim_qk; + params.kv_stride_n = + (kv_layout == QKVLayout::kNHD) ? num_kv_heads * head_dim_vo : head_dim_vo; + params.kv_stride_h = (kv_layout == QKVLayout::kNHD) ? head_dim_vo : kv_len * head_dim_vo; + params.window_left = window_left; + params.kv_chunk_size = 0; + + ADDITIONAL_PARAMS_SETTER + + cudaError_t status = + flashinfer::SingleDecodeWithKVCacheDispatched( + params, static_cast(tmp.data_ptr()), stream); + TORCH_CHECK(status == cudaSuccess, "SingleDecodeWithKVCache kernel launch failed, error: " + + std::string(cudaGetErrorString(status))); + return true; + }); +} diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/single_decode_config.inc b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/single_decode_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..2b1561be467783c431aa53faee1c0bc2e356cf67 --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/single_decode_config.inc @@ -0,0 +1,64 @@ +#pragma once +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta +#define ADDITIONAL_PARAMS_SETTER params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, ...) {\ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ +} + +using namespace flashinfer; + +using DTypeQ = nv_bfloat16; +using DTypeKV = __nv_fp8_e4m3; +using DTypeO = nv_bfloat16; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 128; +constexpr int HEAD_DIM_VO = 128; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = false; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = int32_t; + DTypeQ* q; + DTypeKV* k; + DTypeKV* v; + DTypeO* o; + float* lse; + float* maybe_alibi_slopes; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; + + uint32_t kv_len; + uint32_t num_qo_heads; + uint32_t num_kv_heads; + uint32_t q_stride_n; + uint32_t q_stride_h; + uint32_t kv_stride_n; + uint32_t kv_stride_h; + int32_t window_left; + uint32_t kv_chunk_size; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { return 1; } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return kv_len; + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/single_decode_jit_pybind.cu b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/single_decode_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..f4c360d3e4d283239d95810d4b5aa5bbd05cb953 --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/single_decode_jit_pybind.cu @@ -0,0 +1,27 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "pytorch_extension_utils.h" +#include "single_decode_config.inc" + +void single_decode_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Single-request decode with KV-Cache operator + m.def("run", single_decode_with_kv_cache); +} diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/single_decode_kernel.cu b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/single_decode_kernel.cu new file mode 100644 index 0000000000000000000000000000000000000000..c2ec47236fa827ed1ae4272d6a80a852d73ff5f4 --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/single_decode_kernel.cu @@ -0,0 +1,13 @@ +#include +#include "single_decode_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SingleDecodeWithKVCacheDispatched< + 128, PosEncodingMode::kNone, DefaultAttention, Params>( + Params params, nv_bfloat16* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/single_decode.cu b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/single_decode.cu new file mode 100644 index 0000000000000000000000000000000000000000..a22ebd91727745cd8a85beea286d3c80bef3128c --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/single_decode.cu @@ -0,0 +1,104 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include + +#include "pytorch_extension_utils.h" +#include "single_decode_config.inc" + +namespace flashinfer { + +template +cudaError_t SingleDecodeWithKVCacheDispatched(Params params, typename Params::DTypeO* tmp, + cudaStream_t stream); +} // namespace flashinfer + +using namespace flashinfer; + +void single_decode_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS) { + CHECK_INPUT(q); + CHECK_INPUT(k); + CHECK_INPUT(v); + CHECK_INPUT(tmp); + auto device = q.device(); + CHECK_EQ(k.device(), device); + CHECK_EQ(v.device(), device); + CHECK_EQ(tmp.device(), device); + CHECK_DIM(2, q); + CHECK_DIM(3, k); + CHECK_DIM(3, v); + CHECK_SHAPE(k, v); + CHECK_EQ(q.size(1), k.size(2)); + CHECK_EQ(v.scalar_type(), k.scalar_type()); + unsigned int num_qo_heads = q.size(0); + unsigned int head_dim_qk = q.size(1); + unsigned int head_dim_vo = v.size(2); + unsigned int kv_len, num_kv_heads; + QKVLayout kv_layout = static_cast(layout); + if (kv_layout == QKVLayout::kNHD) { + kv_len = k.size(0); + num_kv_heads = k.size(1); + } else { + num_kv_heads = k.size(0); + kv_len = k.size(1); + } + CHECK_GQA_HEAD_DIVISIBLE(num_qo_heads, num_kv_heads); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(device); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + TORCH_CHECK(head_dim_qk == head_dim_vo, + "CUDA cores template only supports equal head dim for QK and VO, please use tensor " + "cores template for different head dim"); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + Params params; + + params.q = static_cast(q.data_ptr()); + params.k = static_cast(k.data_ptr()); + params.v = static_cast(v.data_ptr()); + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.kv_len = kv_len; + params.num_qo_heads = num_qo_heads; + params.num_kv_heads = num_kv_heads; + params.q_stride_n = num_qo_heads * head_dim_qk; + params.q_stride_h = head_dim_qk; + params.kv_stride_n = + (kv_layout == QKVLayout::kNHD) ? num_kv_heads * head_dim_vo : head_dim_vo; + params.kv_stride_h = (kv_layout == QKVLayout::kNHD) ? head_dim_vo : kv_len * head_dim_vo; + params.window_left = window_left; + params.kv_chunk_size = 0; + + ADDITIONAL_PARAMS_SETTER + + cudaError_t status = + flashinfer::SingleDecodeWithKVCacheDispatched( + params, static_cast(tmp.data_ptr()), stream); + TORCH_CHECK(status == cudaSuccess, "SingleDecodeWithKVCache kernel launch failed, error: " + + std::string(cudaGetErrorString(status))); + return true; + }); +} diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/single_decode_config.inc b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/single_decode_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..38fc89abcabefdfe2db2b542eb549aac407fd7b1 --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/single_decode_config.inc @@ -0,0 +1,64 @@ +#pragma once +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta +#define ADDITIONAL_PARAMS_SETTER params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, ...) {\ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ +} + +using namespace flashinfer; + +using DTypeQ = nv_bfloat16; +using DTypeKV = __nv_fp8_e4m3; +using DTypeO = nv_bfloat16; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 256; +constexpr int HEAD_DIM_VO = 256; +constexpr auto USE_LOGITS_SOFT_CAP = true; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = true; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = int32_t; + DTypeQ* q; + DTypeKV* k; + DTypeKV* v; + DTypeO* o; + float* lse; + float* maybe_alibi_slopes; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; + + uint32_t kv_len; + uint32_t num_qo_heads; + uint32_t num_kv_heads; + uint32_t q_stride_n; + uint32_t q_stride_h; + uint32_t kv_stride_n; + uint32_t kv_stride_h; + int32_t window_left; + uint32_t kv_chunk_size; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { return 1; } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return kv_len; + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/single_decode_jit_pybind.cu b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/single_decode_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..f4c360d3e4d283239d95810d4b5aa5bbd05cb953 --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/single_decode_jit_pybind.cu @@ -0,0 +1,27 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "pytorch_extension_utils.h" +#include "single_decode_config.inc" + +void single_decode_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Single-request decode with KV-Cache operator + m.def("run", single_decode_with_kv_cache); +} diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/single_decode_kernel.cu b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/single_decode_kernel.cu new file mode 100644 index 0000000000000000000000000000000000000000..befd17566ba95c132a8a3fd19a2b756a59ada350 --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/single_decode_kernel.cu @@ -0,0 +1,13 @@ +#include +#include "single_decode_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SingleDecodeWithKVCacheDispatched< + 256, PosEncodingMode::kNone, DefaultAttention, Params>( + Params params, nv_bfloat16* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/single_decode.cu b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/single_decode.cu new file mode 100644 index 0000000000000000000000000000000000000000..a22ebd91727745cd8a85beea286d3c80bef3128c --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/single_decode.cu @@ -0,0 +1,104 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include + +#include "pytorch_extension_utils.h" +#include "single_decode_config.inc" + +namespace flashinfer { + +template +cudaError_t SingleDecodeWithKVCacheDispatched(Params params, typename Params::DTypeO* tmp, + cudaStream_t stream); +} // namespace flashinfer + +using namespace flashinfer; + +void single_decode_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS) { + CHECK_INPUT(q); + CHECK_INPUT(k); + CHECK_INPUT(v); + CHECK_INPUT(tmp); + auto device = q.device(); + CHECK_EQ(k.device(), device); + CHECK_EQ(v.device(), device); + CHECK_EQ(tmp.device(), device); + CHECK_DIM(2, q); + CHECK_DIM(3, k); + CHECK_DIM(3, v); + CHECK_SHAPE(k, v); + CHECK_EQ(q.size(1), k.size(2)); + CHECK_EQ(v.scalar_type(), k.scalar_type()); + unsigned int num_qo_heads = q.size(0); + unsigned int head_dim_qk = q.size(1); + unsigned int head_dim_vo = v.size(2); + unsigned int kv_len, num_kv_heads; + QKVLayout kv_layout = static_cast(layout); + if (kv_layout == QKVLayout::kNHD) { + kv_len = k.size(0); + num_kv_heads = k.size(1); + } else { + num_kv_heads = k.size(0); + kv_len = k.size(1); + } + CHECK_GQA_HEAD_DIVISIBLE(num_qo_heads, num_kv_heads); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(device); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + TORCH_CHECK(head_dim_qk == head_dim_vo, + "CUDA cores template only supports equal head dim for QK and VO, please use tensor " + "cores template for different head dim"); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + Params params; + + params.q = static_cast(q.data_ptr()); + params.k = static_cast(k.data_ptr()); + params.v = static_cast(v.data_ptr()); + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.kv_len = kv_len; + params.num_qo_heads = num_qo_heads; + params.num_kv_heads = num_kv_heads; + params.q_stride_n = num_qo_heads * head_dim_qk; + params.q_stride_h = head_dim_qk; + params.kv_stride_n = + (kv_layout == QKVLayout::kNHD) ? num_kv_heads * head_dim_vo : head_dim_vo; + params.kv_stride_h = (kv_layout == QKVLayout::kNHD) ? head_dim_vo : kv_len * head_dim_vo; + params.window_left = window_left; + params.kv_chunk_size = 0; + + ADDITIONAL_PARAMS_SETTER + + cudaError_t status = + flashinfer::SingleDecodeWithKVCacheDispatched( + params, static_cast(tmp.data_ptr()), stream); + TORCH_CHECK(status == cudaSuccess, "SingleDecodeWithKVCache kernel launch failed, error: " + + std::string(cudaGetErrorString(status))); + return true; + }); +} diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/single_decode_config.inc b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/single_decode_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..ade1409dbe3d770ba783cbcaed0e92adc3b33fe1 --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/single_decode_config.inc @@ -0,0 +1,64 @@ +#pragma once +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta +#define ADDITIONAL_PARAMS_SETTER params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, ...) {\ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ +} + +using namespace flashinfer; + +using DTypeQ = nv_bfloat16; +using DTypeKV = __nv_fp8_e4m3; +using DTypeO = nv_bfloat16; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 64; +constexpr int HEAD_DIM_VO = 64; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = false; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = int32_t; + DTypeQ* q; + DTypeKV* k; + DTypeKV* v; + DTypeO* o; + float* lse; + float* maybe_alibi_slopes; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; + + uint32_t kv_len; + uint32_t num_qo_heads; + uint32_t num_kv_heads; + uint32_t q_stride_n; + uint32_t q_stride_h; + uint32_t kv_stride_n; + uint32_t kv_stride_h; + int32_t window_left; + uint32_t kv_chunk_size; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { return 1; } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return kv_len; + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/single_decode_jit_pybind.cu b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/single_decode_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..f4c360d3e4d283239d95810d4b5aa5bbd05cb953 --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/single_decode_jit_pybind.cu @@ -0,0 +1,27 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "pytorch_extension_utils.h" +#include "single_decode_config.inc" + +void single_decode_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Single-request decode with KV-Cache operator + m.def("run", single_decode_with_kv_cache); +} diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/single_decode_kernel.cu b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/single_decode_kernel.cu new file mode 100644 index 0000000000000000000000000000000000000000..a72a34cca58a670438ef6d99276a121bc4cd0ef5 --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/single_decode_kernel.cu @@ -0,0 +1,13 @@ +#include +#include "single_decode_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SingleDecodeWithKVCacheDispatched< + 64, PosEncodingMode::kNone, DefaultAttention, Params>( + Params params, nv_bfloat16* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/single_decode.cu b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/single_decode.cu new file mode 100644 index 0000000000000000000000000000000000000000..a22ebd91727745cd8a85beea286d3c80bef3128c --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/single_decode.cu @@ -0,0 +1,104 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include + +#include "pytorch_extension_utils.h" +#include "single_decode_config.inc" + +namespace flashinfer { + +template +cudaError_t SingleDecodeWithKVCacheDispatched(Params params, typename Params::DTypeO* tmp, + cudaStream_t stream); +} // namespace flashinfer + +using namespace flashinfer; + +void single_decode_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS) { + CHECK_INPUT(q); + CHECK_INPUT(k); + CHECK_INPUT(v); + CHECK_INPUT(tmp); + auto device = q.device(); + CHECK_EQ(k.device(), device); + CHECK_EQ(v.device(), device); + CHECK_EQ(tmp.device(), device); + CHECK_DIM(2, q); + CHECK_DIM(3, k); + CHECK_DIM(3, v); + CHECK_SHAPE(k, v); + CHECK_EQ(q.size(1), k.size(2)); + CHECK_EQ(v.scalar_type(), k.scalar_type()); + unsigned int num_qo_heads = q.size(0); + unsigned int head_dim_qk = q.size(1); + unsigned int head_dim_vo = v.size(2); + unsigned int kv_len, num_kv_heads; + QKVLayout kv_layout = static_cast(layout); + if (kv_layout == QKVLayout::kNHD) { + kv_len = k.size(0); + num_kv_heads = k.size(1); + } else { + num_kv_heads = k.size(0); + kv_len = k.size(1); + } + CHECK_GQA_HEAD_DIVISIBLE(num_qo_heads, num_kv_heads); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(device); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + TORCH_CHECK(head_dim_qk == head_dim_vo, + "CUDA cores template only supports equal head dim for QK and VO, please use tensor " + "cores template for different head dim"); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + Params params; + + params.q = static_cast(q.data_ptr()); + params.k = static_cast(k.data_ptr()); + params.v = static_cast(v.data_ptr()); + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.kv_len = kv_len; + params.num_qo_heads = num_qo_heads; + params.num_kv_heads = num_kv_heads; + params.q_stride_n = num_qo_heads * head_dim_qk; + params.q_stride_h = head_dim_qk; + params.kv_stride_n = + (kv_layout == QKVLayout::kNHD) ? num_kv_heads * head_dim_vo : head_dim_vo; + params.kv_stride_h = (kv_layout == QKVLayout::kNHD) ? head_dim_vo : kv_len * head_dim_vo; + params.window_left = window_left; + params.kv_chunk_size = 0; + + ADDITIONAL_PARAMS_SETTER + + cudaError_t status = + flashinfer::SingleDecodeWithKVCacheDispatched( + params, static_cast(tmp.data_ptr()), stream); + TORCH_CHECK(status == cudaSuccess, "SingleDecodeWithKVCache kernel launch failed, error: " + + std::string(cudaGetErrorString(status))); + return true; + }); +} diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/single_decode_config.inc b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/single_decode_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..55bf07e597242f00dc78836540e61bbe6476b6c3 --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/single_decode_config.inc @@ -0,0 +1,64 @@ +#pragma once +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta +#define ADDITIONAL_PARAMS_SETTER params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, ...) {\ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ +} + +using namespace flashinfer; + +using DTypeQ = nv_bfloat16; +using DTypeKV = __nv_fp8_e4m3; +using DTypeO = nv_bfloat16; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 64; +constexpr int HEAD_DIM_VO = 64; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = true; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = int32_t; + DTypeQ* q; + DTypeKV* k; + DTypeKV* v; + DTypeO* o; + float* lse; + float* maybe_alibi_slopes; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; + + uint32_t kv_len; + uint32_t num_qo_heads; + uint32_t num_kv_heads; + uint32_t q_stride_n; + uint32_t q_stride_h; + uint32_t kv_stride_n; + uint32_t kv_stride_h; + int32_t window_left; + uint32_t kv_chunk_size; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { return 1; } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return kv_len; + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/single_decode_jit_pybind.cu b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/single_decode_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..f4c360d3e4d283239d95810d4b5aa5bbd05cb953 --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/single_decode_jit_pybind.cu @@ -0,0 +1,27 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "pytorch_extension_utils.h" +#include "single_decode_config.inc" + +void single_decode_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Single-request decode with KV-Cache operator + m.def("run", single_decode_with_kv_cache); +} diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/single_decode_kernel.cu b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/single_decode_kernel.cu new file mode 100644 index 0000000000000000000000000000000000000000..eff8599261be9b92499ce5fdadcd54596fac35cc --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/single_decode_kernel.cu @@ -0,0 +1,13 @@ +#include +#include "single_decode_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SingleDecodeWithKVCacheDispatched< + 64, PosEncodingMode::kNone, DefaultAttention, Params>( + Params params, nv_bfloat16* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/single_decode.cu b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/single_decode.cu new file mode 100644 index 0000000000000000000000000000000000000000..a22ebd91727745cd8a85beea286d3c80bef3128c --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/single_decode.cu @@ -0,0 +1,104 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include + +#include "pytorch_extension_utils.h" +#include "single_decode_config.inc" + +namespace flashinfer { + +template +cudaError_t SingleDecodeWithKVCacheDispatched(Params params, typename Params::DTypeO* tmp, + cudaStream_t stream); +} // namespace flashinfer + +using namespace flashinfer; + +void single_decode_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS) { + CHECK_INPUT(q); + CHECK_INPUT(k); + CHECK_INPUT(v); + CHECK_INPUT(tmp); + auto device = q.device(); + CHECK_EQ(k.device(), device); + CHECK_EQ(v.device(), device); + CHECK_EQ(tmp.device(), device); + CHECK_DIM(2, q); + CHECK_DIM(3, k); + CHECK_DIM(3, v); + CHECK_SHAPE(k, v); + CHECK_EQ(q.size(1), k.size(2)); + CHECK_EQ(v.scalar_type(), k.scalar_type()); + unsigned int num_qo_heads = q.size(0); + unsigned int head_dim_qk = q.size(1); + unsigned int head_dim_vo = v.size(2); + unsigned int kv_len, num_kv_heads; + QKVLayout kv_layout = static_cast(layout); + if (kv_layout == QKVLayout::kNHD) { + kv_len = k.size(0); + num_kv_heads = k.size(1); + } else { + num_kv_heads = k.size(0); + kv_len = k.size(1); + } + CHECK_GQA_HEAD_DIVISIBLE(num_qo_heads, num_kv_heads); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(device); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + TORCH_CHECK(head_dim_qk == head_dim_vo, + "CUDA cores template only supports equal head dim for QK and VO, please use tensor " + "cores template for different head dim"); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + Params params; + + params.q = static_cast(q.data_ptr()); + params.k = static_cast(k.data_ptr()); + params.v = static_cast(v.data_ptr()); + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.kv_len = kv_len; + params.num_qo_heads = num_qo_heads; + params.num_kv_heads = num_kv_heads; + params.q_stride_n = num_qo_heads * head_dim_qk; + params.q_stride_h = head_dim_qk; + params.kv_stride_n = + (kv_layout == QKVLayout::kNHD) ? num_kv_heads * head_dim_vo : head_dim_vo; + params.kv_stride_h = (kv_layout == QKVLayout::kNHD) ? head_dim_vo : kv_len * head_dim_vo; + params.window_left = window_left; + params.kv_chunk_size = 0; + + ADDITIONAL_PARAMS_SETTER + + cudaError_t status = + flashinfer::SingleDecodeWithKVCacheDispatched( + params, static_cast(tmp.data_ptr()), stream); + TORCH_CHECK(status == cudaSuccess, "SingleDecodeWithKVCache kernel launch failed, error: " + + std::string(cudaGetErrorString(status))); + return true; + }); +} diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/single_decode_config.inc b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/single_decode_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..a27d173e49dd1f18827e46b0a594a3a0b56a5624 --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/single_decode_config.inc @@ -0,0 +1,64 @@ +#pragma once +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta +#define ADDITIONAL_PARAMS_SETTER params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, ...) {\ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ +} + +using namespace flashinfer; + +using DTypeQ = half; +using DTypeKV = __nv_fp8_e4m3; +using DTypeO = half; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 128; +constexpr int HEAD_DIM_VO = 128; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = false; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = int32_t; + DTypeQ* q; + DTypeKV* k; + DTypeKV* v; + DTypeO* o; + float* lse; + float* maybe_alibi_slopes; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; + + uint32_t kv_len; + uint32_t num_qo_heads; + uint32_t num_kv_heads; + uint32_t q_stride_n; + uint32_t q_stride_h; + uint32_t kv_stride_n; + uint32_t kv_stride_h; + int32_t window_left; + uint32_t kv_chunk_size; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { return 1; } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return kv_len; + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/single_decode_jit_pybind.cu b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/single_decode_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..f4c360d3e4d283239d95810d4b5aa5bbd05cb953 --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/single_decode_jit_pybind.cu @@ -0,0 +1,27 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "pytorch_extension_utils.h" +#include "single_decode_config.inc" + +void single_decode_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Single-request decode with KV-Cache operator + m.def("run", single_decode_with_kv_cache); +} diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/single_decode_kernel.cu b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/single_decode_kernel.cu new file mode 100644 index 0000000000000000000000000000000000000000..17a0231ddb48012785788dd59a9f46d53b80f5d6 --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/single_decode_kernel.cu @@ -0,0 +1,13 @@ +#include +#include "single_decode_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SingleDecodeWithKVCacheDispatched< + 128, PosEncodingMode::kNone, DefaultAttention, Params>( + Params params, half* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/single_decode.cu b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/single_decode.cu new file mode 100644 index 0000000000000000000000000000000000000000..a22ebd91727745cd8a85beea286d3c80bef3128c --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/single_decode.cu @@ -0,0 +1,104 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include + +#include "pytorch_extension_utils.h" +#include "single_decode_config.inc" + +namespace flashinfer { + +template +cudaError_t SingleDecodeWithKVCacheDispatched(Params params, typename Params::DTypeO* tmp, + cudaStream_t stream); +} // namespace flashinfer + +using namespace flashinfer; + +void single_decode_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS) { + CHECK_INPUT(q); + CHECK_INPUT(k); + CHECK_INPUT(v); + CHECK_INPUT(tmp); + auto device = q.device(); + CHECK_EQ(k.device(), device); + CHECK_EQ(v.device(), device); + CHECK_EQ(tmp.device(), device); + CHECK_DIM(2, q); + CHECK_DIM(3, k); + CHECK_DIM(3, v); + CHECK_SHAPE(k, v); + CHECK_EQ(q.size(1), k.size(2)); + CHECK_EQ(v.scalar_type(), k.scalar_type()); + unsigned int num_qo_heads = q.size(0); + unsigned int head_dim_qk = q.size(1); + unsigned int head_dim_vo = v.size(2); + unsigned int kv_len, num_kv_heads; + QKVLayout kv_layout = static_cast(layout); + if (kv_layout == QKVLayout::kNHD) { + kv_len = k.size(0); + num_kv_heads = k.size(1); + } else { + num_kv_heads = k.size(0); + kv_len = k.size(1); + } + CHECK_GQA_HEAD_DIVISIBLE(num_qo_heads, num_kv_heads); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(device); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + TORCH_CHECK(head_dim_qk == head_dim_vo, + "CUDA cores template only supports equal head dim for QK and VO, please use tensor " + "cores template for different head dim"); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + Params params; + + params.q = static_cast(q.data_ptr()); + params.k = static_cast(k.data_ptr()); + params.v = static_cast(v.data_ptr()); + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.kv_len = kv_len; + params.num_qo_heads = num_qo_heads; + params.num_kv_heads = num_kv_heads; + params.q_stride_n = num_qo_heads * head_dim_qk; + params.q_stride_h = head_dim_qk; + params.kv_stride_n = + (kv_layout == QKVLayout::kNHD) ? num_kv_heads * head_dim_vo : head_dim_vo; + params.kv_stride_h = (kv_layout == QKVLayout::kNHD) ? head_dim_vo : kv_len * head_dim_vo; + params.window_left = window_left; + params.kv_chunk_size = 0; + + ADDITIONAL_PARAMS_SETTER + + cudaError_t status = + flashinfer::SingleDecodeWithKVCacheDispatched( + params, static_cast(tmp.data_ptr()), stream); + TORCH_CHECK(status == cudaSuccess, "SingleDecodeWithKVCache kernel launch failed, error: " + + std::string(cudaGetErrorString(status))); + return true; + }); +} diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/single_decode_config.inc b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/single_decode_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..dd73c1566f4a1d620aa66ed1cbc549f7936b5732 --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/single_decode_config.inc @@ -0,0 +1,64 @@ +#pragma once +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta +#define ADDITIONAL_PARAMS_SETTER params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, ...) {\ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ +} + +using namespace flashinfer; + +using DTypeQ = half; +using DTypeKV = __nv_fp8_e4m3; +using DTypeO = half; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 256; +constexpr int HEAD_DIM_VO = 256; +constexpr auto USE_LOGITS_SOFT_CAP = true; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = true; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = int32_t; + DTypeQ* q; + DTypeKV* k; + DTypeKV* v; + DTypeO* o; + float* lse; + float* maybe_alibi_slopes; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; + + uint32_t kv_len; + uint32_t num_qo_heads; + uint32_t num_kv_heads; + uint32_t q_stride_n; + uint32_t q_stride_h; + uint32_t kv_stride_n; + uint32_t kv_stride_h; + int32_t window_left; + uint32_t kv_chunk_size; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { return 1; } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return kv_len; + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/single_decode_jit_pybind.cu b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/single_decode_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..f4c360d3e4d283239d95810d4b5aa5bbd05cb953 --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/single_decode_jit_pybind.cu @@ -0,0 +1,27 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "pytorch_extension_utils.h" +#include "single_decode_config.inc" + +void single_decode_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Single-request decode with KV-Cache operator + m.def("run", single_decode_with_kv_cache); +} diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/single_decode_kernel.cu b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/single_decode_kernel.cu new file mode 100644 index 0000000000000000000000000000000000000000..a93004e275d20863b19fa6de83f5f531bb5a2de2 --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/single_decode_kernel.cu @@ -0,0 +1,13 @@ +#include +#include "single_decode_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SingleDecodeWithKVCacheDispatched< + 256, PosEncodingMode::kNone, DefaultAttention, Params>( + Params params, half* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/single_decode.cu b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/single_decode.cu new file mode 100644 index 0000000000000000000000000000000000000000..a22ebd91727745cd8a85beea286d3c80bef3128c --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/single_decode.cu @@ -0,0 +1,104 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include + +#include "pytorch_extension_utils.h" +#include "single_decode_config.inc" + +namespace flashinfer { + +template +cudaError_t SingleDecodeWithKVCacheDispatched(Params params, typename Params::DTypeO* tmp, + cudaStream_t stream); +} // namespace flashinfer + +using namespace flashinfer; + +void single_decode_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS) { + CHECK_INPUT(q); + CHECK_INPUT(k); + CHECK_INPUT(v); + CHECK_INPUT(tmp); + auto device = q.device(); + CHECK_EQ(k.device(), device); + CHECK_EQ(v.device(), device); + CHECK_EQ(tmp.device(), device); + CHECK_DIM(2, q); + CHECK_DIM(3, k); + CHECK_DIM(3, v); + CHECK_SHAPE(k, v); + CHECK_EQ(q.size(1), k.size(2)); + CHECK_EQ(v.scalar_type(), k.scalar_type()); + unsigned int num_qo_heads = q.size(0); + unsigned int head_dim_qk = q.size(1); + unsigned int head_dim_vo = v.size(2); + unsigned int kv_len, num_kv_heads; + QKVLayout kv_layout = static_cast(layout); + if (kv_layout == QKVLayout::kNHD) { + kv_len = k.size(0); + num_kv_heads = k.size(1); + } else { + num_kv_heads = k.size(0); + kv_len = k.size(1); + } + CHECK_GQA_HEAD_DIVISIBLE(num_qo_heads, num_kv_heads); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(device); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + TORCH_CHECK(head_dim_qk == head_dim_vo, + "CUDA cores template only supports equal head dim for QK and VO, please use tensor " + "cores template for different head dim"); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + Params params; + + params.q = static_cast(q.data_ptr()); + params.k = static_cast(k.data_ptr()); + params.v = static_cast(v.data_ptr()); + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.kv_len = kv_len; + params.num_qo_heads = num_qo_heads; + params.num_kv_heads = num_kv_heads; + params.q_stride_n = num_qo_heads * head_dim_qk; + params.q_stride_h = head_dim_qk; + params.kv_stride_n = + (kv_layout == QKVLayout::kNHD) ? num_kv_heads * head_dim_vo : head_dim_vo; + params.kv_stride_h = (kv_layout == QKVLayout::kNHD) ? head_dim_vo : kv_len * head_dim_vo; + params.window_left = window_left; + params.kv_chunk_size = 0; + + ADDITIONAL_PARAMS_SETTER + + cudaError_t status = + flashinfer::SingleDecodeWithKVCacheDispatched( + params, static_cast(tmp.data_ptr()), stream); + TORCH_CHECK(status == cudaSuccess, "SingleDecodeWithKVCache kernel launch failed, error: " + + std::string(cudaGetErrorString(status))); + return true; + }); +} diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/single_decode_config.inc b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/single_decode_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..36df6ea7e0e0bee2a12fab4816a20f74eee4f206 --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/single_decode_config.inc @@ -0,0 +1,64 @@ +#pragma once +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta +#define ADDITIONAL_PARAMS_SETTER params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, ...) {\ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ +} + +using namespace flashinfer; + +using DTypeQ = half; +using DTypeKV = __nv_fp8_e4m3; +using DTypeO = half; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 64; +constexpr int HEAD_DIM_VO = 64; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = false; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = int32_t; + DTypeQ* q; + DTypeKV* k; + DTypeKV* v; + DTypeO* o; + float* lse; + float* maybe_alibi_slopes; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; + + uint32_t kv_len; + uint32_t num_qo_heads; + uint32_t num_kv_heads; + uint32_t q_stride_n; + uint32_t q_stride_h; + uint32_t kv_stride_n; + uint32_t kv_stride_h; + int32_t window_left; + uint32_t kv_chunk_size; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { return 1; } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return kv_len; + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/single_decode_jit_pybind.cu b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/single_decode_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..f4c360d3e4d283239d95810d4b5aa5bbd05cb953 --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/single_decode_jit_pybind.cu @@ -0,0 +1,27 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "pytorch_extension_utils.h" +#include "single_decode_config.inc" + +void single_decode_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Single-request decode with KV-Cache operator + m.def("run", single_decode_with_kv_cache); +} diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/single_decode_kernel.cu b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/single_decode_kernel.cu new file mode 100644 index 0000000000000000000000000000000000000000..b9f763d9180c60024a2e520149d3a200c242dbac --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/single_decode_kernel.cu @@ -0,0 +1,13 @@ +#include +#include "single_decode_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SingleDecodeWithKVCacheDispatched< + 64, PosEncodingMode::kNone, DefaultAttention, Params>( + Params params, half* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/single_decode.cu b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/single_decode.cu new file mode 100644 index 0000000000000000000000000000000000000000..a22ebd91727745cd8a85beea286d3c80bef3128c --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/single_decode.cu @@ -0,0 +1,104 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include + +#include "pytorch_extension_utils.h" +#include "single_decode_config.inc" + +namespace flashinfer { + +template +cudaError_t SingleDecodeWithKVCacheDispatched(Params params, typename Params::DTypeO* tmp, + cudaStream_t stream); +} // namespace flashinfer + +using namespace flashinfer; + +void single_decode_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS) { + CHECK_INPUT(q); + CHECK_INPUT(k); + CHECK_INPUT(v); + CHECK_INPUT(tmp); + auto device = q.device(); + CHECK_EQ(k.device(), device); + CHECK_EQ(v.device(), device); + CHECK_EQ(tmp.device(), device); + CHECK_DIM(2, q); + CHECK_DIM(3, k); + CHECK_DIM(3, v); + CHECK_SHAPE(k, v); + CHECK_EQ(q.size(1), k.size(2)); + CHECK_EQ(v.scalar_type(), k.scalar_type()); + unsigned int num_qo_heads = q.size(0); + unsigned int head_dim_qk = q.size(1); + unsigned int head_dim_vo = v.size(2); + unsigned int kv_len, num_kv_heads; + QKVLayout kv_layout = static_cast(layout); + if (kv_layout == QKVLayout::kNHD) { + kv_len = k.size(0); + num_kv_heads = k.size(1); + } else { + num_kv_heads = k.size(0); + kv_len = k.size(1); + } + CHECK_GQA_HEAD_DIVISIBLE(num_qo_heads, num_kv_heads); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(device); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + TORCH_CHECK(head_dim_qk == head_dim_vo, + "CUDA cores template only supports equal head dim for QK and VO, please use tensor " + "cores template for different head dim"); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + Params params; + + params.q = static_cast(q.data_ptr()); + params.k = static_cast(k.data_ptr()); + params.v = static_cast(v.data_ptr()); + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.kv_len = kv_len; + params.num_qo_heads = num_qo_heads; + params.num_kv_heads = num_kv_heads; + params.q_stride_n = num_qo_heads * head_dim_qk; + params.q_stride_h = head_dim_qk; + params.kv_stride_n = + (kv_layout == QKVLayout::kNHD) ? num_kv_heads * head_dim_vo : head_dim_vo; + params.kv_stride_h = (kv_layout == QKVLayout::kNHD) ? head_dim_vo : kv_len * head_dim_vo; + params.window_left = window_left; + params.kv_chunk_size = 0; + + ADDITIONAL_PARAMS_SETTER + + cudaError_t status = + flashinfer::SingleDecodeWithKVCacheDispatched( + params, static_cast(tmp.data_ptr()), stream); + TORCH_CHECK(status == cudaSuccess, "SingleDecodeWithKVCache kernel launch failed, error: " + + std::string(cudaGetErrorString(status))); + return true; + }); +} diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/single_decode_config.inc b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/single_decode_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..c0fbe1ba1158e548072c2572f50e91a6fb522ceb --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/single_decode_config.inc @@ -0,0 +1,64 @@ +#pragma once +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta +#define ADDITIONAL_PARAMS_SETTER params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, ...) {\ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ +} + +using namespace flashinfer; + +using DTypeQ = half; +using DTypeKV = __nv_fp8_e4m3; +using DTypeO = half; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 64; +constexpr int HEAD_DIM_VO = 64; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = true; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = int32_t; + DTypeQ* q; + DTypeKV* k; + DTypeKV* v; + DTypeO* o; + float* lse; + float* maybe_alibi_slopes; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; + + uint32_t kv_len; + uint32_t num_qo_heads; + uint32_t num_kv_heads; + uint32_t q_stride_n; + uint32_t q_stride_h; + uint32_t kv_stride_n; + uint32_t kv_stride_h; + int32_t window_left; + uint32_t kv_chunk_size; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { return 1; } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return kv_len; + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/single_decode_jit_pybind.cu b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/single_decode_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..f4c360d3e4d283239d95810d4b5aa5bbd05cb953 --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/single_decode_jit_pybind.cu @@ -0,0 +1,27 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "pytorch_extension_utils.h" +#include "single_decode_config.inc" + +void single_decode_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Single-request decode with KV-Cache operator + m.def("run", single_decode_with_kv_cache); +} diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/single_decode_kernel.cu b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/single_decode_kernel.cu new file mode 100644 index 0000000000000000000000000000000000000000..2ef7b03887b2006b4b3b265dacbbd59367a37664 --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/single_decode_kernel.cu @@ -0,0 +1,13 @@ +#include +#include "single_decode_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SingleDecodeWithKVCacheDispatched< + 64, PosEncodingMode::kNone, DefaultAttention, Params>( + Params params, half* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/single_decode.cu b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/single_decode.cu new file mode 100644 index 0000000000000000000000000000000000000000..a22ebd91727745cd8a85beea286d3c80bef3128c --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/single_decode.cu @@ -0,0 +1,104 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include + +#include "pytorch_extension_utils.h" +#include "single_decode_config.inc" + +namespace flashinfer { + +template +cudaError_t SingleDecodeWithKVCacheDispatched(Params params, typename Params::DTypeO* tmp, + cudaStream_t stream); +} // namespace flashinfer + +using namespace flashinfer; + +void single_decode_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS) { + CHECK_INPUT(q); + CHECK_INPUT(k); + CHECK_INPUT(v); + CHECK_INPUT(tmp); + auto device = q.device(); + CHECK_EQ(k.device(), device); + CHECK_EQ(v.device(), device); + CHECK_EQ(tmp.device(), device); + CHECK_DIM(2, q); + CHECK_DIM(3, k); + CHECK_DIM(3, v); + CHECK_SHAPE(k, v); + CHECK_EQ(q.size(1), k.size(2)); + CHECK_EQ(v.scalar_type(), k.scalar_type()); + unsigned int num_qo_heads = q.size(0); + unsigned int head_dim_qk = q.size(1); + unsigned int head_dim_vo = v.size(2); + unsigned int kv_len, num_kv_heads; + QKVLayout kv_layout = static_cast(layout); + if (kv_layout == QKVLayout::kNHD) { + kv_len = k.size(0); + num_kv_heads = k.size(1); + } else { + num_kv_heads = k.size(0); + kv_len = k.size(1); + } + CHECK_GQA_HEAD_DIVISIBLE(num_qo_heads, num_kv_heads); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(device); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + TORCH_CHECK(head_dim_qk == head_dim_vo, + "CUDA cores template only supports equal head dim for QK and VO, please use tensor " + "cores template for different head dim"); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + Params params; + + params.q = static_cast(q.data_ptr()); + params.k = static_cast(k.data_ptr()); + params.v = static_cast(v.data_ptr()); + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.kv_len = kv_len; + params.num_qo_heads = num_qo_heads; + params.num_kv_heads = num_kv_heads; + params.q_stride_n = num_qo_heads * head_dim_qk; + params.q_stride_h = head_dim_qk; + params.kv_stride_n = + (kv_layout == QKVLayout::kNHD) ? num_kv_heads * head_dim_vo : head_dim_vo; + params.kv_stride_h = (kv_layout == QKVLayout::kNHD) ? head_dim_vo : kv_len * head_dim_vo; + params.window_left = window_left; + params.kv_chunk_size = 0; + + ADDITIONAL_PARAMS_SETTER + + cudaError_t status = + flashinfer::SingleDecodeWithKVCacheDispatched( + params, static_cast(tmp.data_ptr()), stream); + TORCH_CHECK(status == cudaSuccess, "SingleDecodeWithKVCache kernel launch failed, error: " + + std::string(cudaGetErrorString(status))); + return true; + }); +} diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/single_decode_config.inc b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/single_decode_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..47acdabf6bce5b67df6b6ffd518fdaa1d1df4fd3 --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/single_decode_config.inc @@ -0,0 +1,64 @@ +#pragma once +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta +#define ADDITIONAL_PARAMS_SETTER params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, ...) {\ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ +} + +using namespace flashinfer; + +using DTypeQ = half; +using DTypeKV = half; +using DTypeO = half; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 128; +constexpr int HEAD_DIM_VO = 128; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = false; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = int32_t; + DTypeQ* q; + DTypeKV* k; + DTypeKV* v; + DTypeO* o; + float* lse; + float* maybe_alibi_slopes; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; + + uint32_t kv_len; + uint32_t num_qo_heads; + uint32_t num_kv_heads; + uint32_t q_stride_n; + uint32_t q_stride_h; + uint32_t kv_stride_n; + uint32_t kv_stride_h; + int32_t window_left; + uint32_t kv_chunk_size; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { return 1; } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return kv_len; + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/single_decode_jit_pybind.cu b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/single_decode_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..f4c360d3e4d283239d95810d4b5aa5bbd05cb953 --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/single_decode_jit_pybind.cu @@ -0,0 +1,27 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "pytorch_extension_utils.h" +#include "single_decode_config.inc" + +void single_decode_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Single-request decode with KV-Cache operator + m.def("run", single_decode_with_kv_cache); +} diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/single_decode_kernel.cu b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/single_decode_kernel.cu new file mode 100644 index 0000000000000000000000000000000000000000..17a0231ddb48012785788dd59a9f46d53b80f5d6 --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False/single_decode_kernel.cu @@ -0,0 +1,13 @@ +#include +#include "single_decode_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SingleDecodeWithKVCacheDispatched< + 128, PosEncodingMode::kNone, DefaultAttention, Params>( + Params params, half* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/single_decode.cu b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/single_decode.cu new file mode 100644 index 0000000000000000000000000000000000000000..a22ebd91727745cd8a85beea286d3c80bef3128c --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/single_decode.cu @@ -0,0 +1,104 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include + +#include "pytorch_extension_utils.h" +#include "single_decode_config.inc" + +namespace flashinfer { + +template +cudaError_t SingleDecodeWithKVCacheDispatched(Params params, typename Params::DTypeO* tmp, + cudaStream_t stream); +} // namespace flashinfer + +using namespace flashinfer; + +void single_decode_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS) { + CHECK_INPUT(q); + CHECK_INPUT(k); + CHECK_INPUT(v); + CHECK_INPUT(tmp); + auto device = q.device(); + CHECK_EQ(k.device(), device); + CHECK_EQ(v.device(), device); + CHECK_EQ(tmp.device(), device); + CHECK_DIM(2, q); + CHECK_DIM(3, k); + CHECK_DIM(3, v); + CHECK_SHAPE(k, v); + CHECK_EQ(q.size(1), k.size(2)); + CHECK_EQ(v.scalar_type(), k.scalar_type()); + unsigned int num_qo_heads = q.size(0); + unsigned int head_dim_qk = q.size(1); + unsigned int head_dim_vo = v.size(2); + unsigned int kv_len, num_kv_heads; + QKVLayout kv_layout = static_cast(layout); + if (kv_layout == QKVLayout::kNHD) { + kv_len = k.size(0); + num_kv_heads = k.size(1); + } else { + num_kv_heads = k.size(0); + kv_len = k.size(1); + } + CHECK_GQA_HEAD_DIVISIBLE(num_qo_heads, num_kv_heads); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(device); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + TORCH_CHECK(head_dim_qk == head_dim_vo, + "CUDA cores template only supports equal head dim for QK and VO, please use tensor " + "cores template for different head dim"); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + Params params; + + params.q = static_cast(q.data_ptr()); + params.k = static_cast(k.data_ptr()); + params.v = static_cast(v.data_ptr()); + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.kv_len = kv_len; + params.num_qo_heads = num_qo_heads; + params.num_kv_heads = num_kv_heads; + params.q_stride_n = num_qo_heads * head_dim_qk; + params.q_stride_h = head_dim_qk; + params.kv_stride_n = + (kv_layout == QKVLayout::kNHD) ? num_kv_heads * head_dim_vo : head_dim_vo; + params.kv_stride_h = (kv_layout == QKVLayout::kNHD) ? head_dim_vo : kv_len * head_dim_vo; + params.window_left = window_left; + params.kv_chunk_size = 0; + + ADDITIONAL_PARAMS_SETTER + + cudaError_t status = + flashinfer::SingleDecodeWithKVCacheDispatched( + params, static_cast(tmp.data_ptr()), stream); + TORCH_CHECK(status == cudaSuccess, "SingleDecodeWithKVCache kernel launch failed, error: " + + std::string(cudaGetErrorString(status))); + return true; + }); +} diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/single_decode_config.inc b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/single_decode_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..262cbe9de74a1f674efa2264db0637ed3b1de466 --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/single_decode_config.inc @@ -0,0 +1,64 @@ +#pragma once +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta +#define ADDITIONAL_PARAMS_SETTER params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, ...) {\ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ +} + +using namespace flashinfer; + +using DTypeQ = half; +using DTypeKV = half; +using DTypeO = half; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 256; +constexpr int HEAD_DIM_VO = 256; +constexpr auto USE_LOGITS_SOFT_CAP = true; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = true; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = int32_t; + DTypeQ* q; + DTypeKV* k; + DTypeKV* v; + DTypeO* o; + float* lse; + float* maybe_alibi_slopes; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; + + uint32_t kv_len; + uint32_t num_qo_heads; + uint32_t num_kv_heads; + uint32_t q_stride_n; + uint32_t q_stride_h; + uint32_t kv_stride_n; + uint32_t kv_stride_h; + int32_t window_left; + uint32_t kv_chunk_size; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { return 1; } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return kv_len; + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/single_decode_jit_pybind.cu b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/single_decode_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..f4c360d3e4d283239d95810d4b5aa5bbd05cb953 --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/single_decode_jit_pybind.cu @@ -0,0 +1,27 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "pytorch_extension_utils.h" +#include "single_decode_config.inc" + +void single_decode_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Single-request decode with KV-Cache operator + m.def("run", single_decode_with_kv_cache); +} diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/single_decode_kernel.cu b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/single_decode_kernel.cu new file mode 100644 index 0000000000000000000000000000000000000000..a93004e275d20863b19fa6de83f5f531bb5a2de2 --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True/single_decode_kernel.cu @@ -0,0 +1,13 @@ +#include +#include "single_decode_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SingleDecodeWithKVCacheDispatched< + 256, PosEncodingMode::kNone, DefaultAttention, Params>( + Params params, half* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/single_decode.cu b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/single_decode.cu new file mode 100644 index 0000000000000000000000000000000000000000..a22ebd91727745cd8a85beea286d3c80bef3128c --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/single_decode.cu @@ -0,0 +1,104 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include + +#include "pytorch_extension_utils.h" +#include "single_decode_config.inc" + +namespace flashinfer { + +template +cudaError_t SingleDecodeWithKVCacheDispatched(Params params, typename Params::DTypeO* tmp, + cudaStream_t stream); +} // namespace flashinfer + +using namespace flashinfer; + +void single_decode_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS) { + CHECK_INPUT(q); + CHECK_INPUT(k); + CHECK_INPUT(v); + CHECK_INPUT(tmp); + auto device = q.device(); + CHECK_EQ(k.device(), device); + CHECK_EQ(v.device(), device); + CHECK_EQ(tmp.device(), device); + CHECK_DIM(2, q); + CHECK_DIM(3, k); + CHECK_DIM(3, v); + CHECK_SHAPE(k, v); + CHECK_EQ(q.size(1), k.size(2)); + CHECK_EQ(v.scalar_type(), k.scalar_type()); + unsigned int num_qo_heads = q.size(0); + unsigned int head_dim_qk = q.size(1); + unsigned int head_dim_vo = v.size(2); + unsigned int kv_len, num_kv_heads; + QKVLayout kv_layout = static_cast(layout); + if (kv_layout == QKVLayout::kNHD) { + kv_len = k.size(0); + num_kv_heads = k.size(1); + } else { + num_kv_heads = k.size(0); + kv_len = k.size(1); + } + CHECK_GQA_HEAD_DIVISIBLE(num_qo_heads, num_kv_heads); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(device); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + TORCH_CHECK(head_dim_qk == head_dim_vo, + "CUDA cores template only supports equal head dim for QK and VO, please use tensor " + "cores template for different head dim"); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + Params params; + + params.q = static_cast(q.data_ptr()); + params.k = static_cast(k.data_ptr()); + params.v = static_cast(v.data_ptr()); + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.kv_len = kv_len; + params.num_qo_heads = num_qo_heads; + params.num_kv_heads = num_kv_heads; + params.q_stride_n = num_qo_heads * head_dim_qk; + params.q_stride_h = head_dim_qk; + params.kv_stride_n = + (kv_layout == QKVLayout::kNHD) ? num_kv_heads * head_dim_vo : head_dim_vo; + params.kv_stride_h = (kv_layout == QKVLayout::kNHD) ? head_dim_vo : kv_len * head_dim_vo; + params.window_left = window_left; + params.kv_chunk_size = 0; + + ADDITIONAL_PARAMS_SETTER + + cudaError_t status = + flashinfer::SingleDecodeWithKVCacheDispatched( + params, static_cast(tmp.data_ptr()), stream); + TORCH_CHECK(status == cudaSuccess, "SingleDecodeWithKVCache kernel launch failed, error: " + + std::string(cudaGetErrorString(status))); + return true; + }); +} diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/single_decode_config.inc b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/single_decode_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..725a4f323da55b56912138cf963235040600c269 --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/single_decode_config.inc @@ -0,0 +1,64 @@ +#pragma once +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta +#define ADDITIONAL_PARAMS_SETTER params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, ...) {\ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ +} + +using namespace flashinfer; + +using DTypeQ = half; +using DTypeKV = half; +using DTypeO = half; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 64; +constexpr int HEAD_DIM_VO = 64; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = false; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = int32_t; + DTypeQ* q; + DTypeKV* k; + DTypeKV* v; + DTypeO* o; + float* lse; + float* maybe_alibi_slopes; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; + + uint32_t kv_len; + uint32_t num_qo_heads; + uint32_t num_kv_heads; + uint32_t q_stride_n; + uint32_t q_stride_h; + uint32_t kv_stride_n; + uint32_t kv_stride_h; + int32_t window_left; + uint32_t kv_chunk_size; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { return 1; } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return kv_len; + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/single_decode_jit_pybind.cu b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/single_decode_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..f4c360d3e4d283239d95810d4b5aa5bbd05cb953 --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/single_decode_jit_pybind.cu @@ -0,0 +1,27 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "pytorch_extension_utils.h" +#include "single_decode_config.inc" + +void single_decode_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Single-request decode with KV-Cache operator + m.def("run", single_decode_with_kv_cache); +} diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/single_decode_kernel.cu b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/single_decode_kernel.cu new file mode 100644 index 0000000000000000000000000000000000000000..b9f763d9180c60024a2e520149d3a200c242dbac --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False/single_decode_kernel.cu @@ -0,0 +1,13 @@ +#include +#include "single_decode_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SingleDecodeWithKVCacheDispatched< + 64, PosEncodingMode::kNone, DefaultAttention, Params>( + Params params, half* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/single_decode.cu b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/single_decode.cu new file mode 100644 index 0000000000000000000000000000000000000000..a22ebd91727745cd8a85beea286d3c80bef3128c --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/single_decode.cu @@ -0,0 +1,104 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include + +#include "pytorch_extension_utils.h" +#include "single_decode_config.inc" + +namespace flashinfer { + +template +cudaError_t SingleDecodeWithKVCacheDispatched(Params params, typename Params::DTypeO* tmp, + cudaStream_t stream); +} // namespace flashinfer + +using namespace flashinfer; + +void single_decode_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS) { + CHECK_INPUT(q); + CHECK_INPUT(k); + CHECK_INPUT(v); + CHECK_INPUT(tmp); + auto device = q.device(); + CHECK_EQ(k.device(), device); + CHECK_EQ(v.device(), device); + CHECK_EQ(tmp.device(), device); + CHECK_DIM(2, q); + CHECK_DIM(3, k); + CHECK_DIM(3, v); + CHECK_SHAPE(k, v); + CHECK_EQ(q.size(1), k.size(2)); + CHECK_EQ(v.scalar_type(), k.scalar_type()); + unsigned int num_qo_heads = q.size(0); + unsigned int head_dim_qk = q.size(1); + unsigned int head_dim_vo = v.size(2); + unsigned int kv_len, num_kv_heads; + QKVLayout kv_layout = static_cast(layout); + if (kv_layout == QKVLayout::kNHD) { + kv_len = k.size(0); + num_kv_heads = k.size(1); + } else { + num_kv_heads = k.size(0); + kv_len = k.size(1); + } + CHECK_GQA_HEAD_DIVISIBLE(num_qo_heads, num_kv_heads); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(device); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + TORCH_CHECK(head_dim_qk == head_dim_vo, + "CUDA cores template only supports equal head dim for QK and VO, please use tensor " + "cores template for different head dim"); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + Params params; + + params.q = static_cast(q.data_ptr()); + params.k = static_cast(k.data_ptr()); + params.v = static_cast(v.data_ptr()); + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.kv_len = kv_len; + params.num_qo_heads = num_qo_heads; + params.num_kv_heads = num_kv_heads; + params.q_stride_n = num_qo_heads * head_dim_qk; + params.q_stride_h = head_dim_qk; + params.kv_stride_n = + (kv_layout == QKVLayout::kNHD) ? num_kv_heads * head_dim_vo : head_dim_vo; + params.kv_stride_h = (kv_layout == QKVLayout::kNHD) ? head_dim_vo : kv_len * head_dim_vo; + params.window_left = window_left; + params.kv_chunk_size = 0; + + ADDITIONAL_PARAMS_SETTER + + cudaError_t status = + flashinfer::SingleDecodeWithKVCacheDispatched( + params, static_cast(tmp.data_ptr()), stream); + TORCH_CHECK(status == cudaSuccess, "SingleDecodeWithKVCache kernel launch failed, error: " + + std::string(cudaGetErrorString(status))); + return true; + }); +} diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/single_decode_config.inc b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/single_decode_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..11e6e7601f827a379001ba92ec19264136689212 --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/single_decode_config.inc @@ -0,0 +1,64 @@ +#pragma once +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta +#define ADDITIONAL_PARAMS_SETTER params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, ...) {\ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ +} + +using namespace flashinfer; + +using DTypeQ = half; +using DTypeKV = half; +using DTypeO = half; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 64; +constexpr int HEAD_DIM_VO = 64; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = true; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = int32_t; + DTypeQ* q; + DTypeKV* k; + DTypeKV* v; + DTypeO* o; + float* lse; + float* maybe_alibi_slopes; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; + + uint32_t kv_len; + uint32_t num_qo_heads; + uint32_t num_kv_heads; + uint32_t q_stride_n; + uint32_t q_stride_h; + uint32_t kv_stride_n; + uint32_t kv_stride_h; + int32_t window_left; + uint32_t kv_chunk_size; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { return 1; } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return kv_len; + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/single_decode_jit_pybind.cu b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/single_decode_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..f4c360d3e4d283239d95810d4b5aa5bbd05cb953 --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/single_decode_jit_pybind.cu @@ -0,0 +1,27 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "pytorch_extension_utils.h" +#include "single_decode_config.inc" + +void single_decode_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Single-request decode with KV-Cache operator + m.def("run", single_decode_with_kv_cache); +} diff --git a/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/single_decode_kernel.cu b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/single_decode_kernel.cu new file mode 100644 index 0000000000000000000000000000000000000000..2ef7b03887b2006b4b3b265dacbbd59367a37664 --- /dev/null +++ b/csrc/generated/single_decode_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False/single_decode_kernel.cu @@ -0,0 +1,13 @@ +#include +#include "single_decode_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SingleDecodeWithKVCacheDispatched< + 64, PosEncodingMode::kNone, DefaultAttention, Params>( + Params params, half* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill.cu new file mode 100644 index 0000000000000000000000000000000000000000..8d77f8420680e70a7929d8b16bd95f5587a01752 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill.cu @@ -0,0 +1,114 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include + +#include "flashinfer/fastdiv.cuh" +#include "pytorch_extension_utils.h" +#include "single_prefill_config.inc" + +namespace flashinfer { + +template +cudaError_t SinglePrefillWithKVCacheDispatched(Params params, typename Params::DTypeO* tmp, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +void single_prefill_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS) { + auto device = q.device(); + unsigned int head_dim_qk = q.size(2); + unsigned int kv_len, qo_len, num_kv_heads, num_qo_heads; + QKVLayout kv_layout = static_cast(layout); + qo_len = q.size(0); + num_qo_heads = q.size(1); + uint32_t q_stride_n = q.stride(0), q_stride_h = q.stride(1), k_stride_n, k_stride_h, v_stride_n, + v_stride_h; + if (kv_layout == QKVLayout::kNHD) { + kv_len = k.size(0); + num_kv_heads = k.size(1); + k_stride_n = k.stride(0); + k_stride_h = k.stride(1); + v_stride_n = v.stride(0); + v_stride_h = v.stride(1); + } else { + kv_len = k.size(1); + num_kv_heads = k.size(0); + k_stride_h = k.stride(0); + k_stride_n = k.stride(1); + v_stride_h = v.stride(0); + v_stride_n = v.stride(1); + } + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == qo_len, lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == num_qo_heads, lse.size(1), q.size(1)); + } + + const MaskMode mask_mode = static_cast(mask_mode_code); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(device); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, Params, + [&] { + Params params; + + params.q = static_cast(q.data_ptr()); + params.k = static_cast(k.data_ptr()); + params.v = static_cast(v.data_ptr()); + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.num_qo_heads = num_qo_heads; + params.num_kv_heads = num_kv_heads; + params.group_size = uint_fastdiv(num_qo_heads / num_kv_heads); + params.qo_len = qo_len; + params.kv_len = kv_len; + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.k_stride_n = k_stride_n; + params.k_stride_h = k_stride_h; + params.v_stride_n = v_stride_n; + params.v_stride_h = v_stride_h; + + params.window_left = window_left; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + cudaError_t status = flashinfer::SinglePrefillWithKVCacheDispatched< + HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + /*use_fp16_qk_reduction=*/USE_FP16_QK_REDUCTION, MASK_MODE, AttentionVariant>( + params, static_cast(tmp.data_ptr()), stream); + TORCH_CHECK(status == cudaSuccess, + "SinglePrefillWithKVCache kernel launch failed, error: " + + std::string(cudaGetErrorString(status))); + return true; + }); +} diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_config.inc b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..a8406e0a6de8995d5242baa051c8549c13356a39 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_config.inc @@ -0,0 +1,83 @@ +#pragma once +#include +#include +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_custom_mask, std::optional maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta +#define ADDITIONAL_PARAMS_SETTER params.maybe_custom_mask = maybe_custom_mask ? static_cast(maybe_custom_mask->data_ptr()): nullptr; \ +params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; + + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, Params, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { \ + constexpr bool use_custom_mask = MASK_MODE == MaskMode::kCustom; \ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ + }) + + +using namespace flashinfer; + +using DTypeQ = nv_bfloat16; +using DTypeKV = nv_bfloat16; +using DTypeO = nv_bfloat16; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 128; +constexpr int HEAD_DIM_VO = 128; +constexpr bool USE_FP16_QK_REDUCTION = false; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = false; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = int32_t; + DTypeQ* q; + DTypeKV* k; + DTypeKV* v; + DTypeO* o; + float* lse; + uint_fastdiv group_size; + + uint8_t* maybe_custom_mask; +float* maybe_alibi_slopes; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; + + + uint32_t qo_len; + uint32_t kv_len; + uint32_t num_qo_heads; + uint32_t num_kv_heads; + uint32_t q_stride_n; + uint32_t q_stride_h; + uint32_t k_stride_n; + uint32_t k_stride_h; + uint32_t v_stride_n; + uint32_t v_stride_h; + uint32_t head_dim; + int32_t window_left; + + bool partition_kv; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { + return qo_len; + } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return kv_len; + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_jit_pybind.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..aaccbc045f9fd8566654a106379c99a053489fc7 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_jit_pybind.cu @@ -0,0 +1,27 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "pytorch_extension_utils.h" +#include "single_prefill_config.inc" + +void single_prefill_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Single-request prefill attention with KV-Cache operator + m.def("run", single_prefill_with_kv_cache); +} diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_0.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..96c9a8e7cf1247b5c741af94c480cf5068a2f97a --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_0.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kNone == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 128, 128, PosEncodingMode::kNone, false, MaskMode::kNone, DefaultAttention, Params>( + Params params, nv_bfloat16* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_1.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..8c4bdb83a9b4e27744f1ca7b8a3b32a72f1bb5f5 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_1.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCausal == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 128, 128, PosEncodingMode::kNone, false, MaskMode::kCausal, DefaultAttention, Params>( + Params params, nv_bfloat16* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_2.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..8d038b88e1988b2313593ad20d0f8fcc108ce129 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_2.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCustom == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 128, 128, PosEncodingMode::kNone, false, MaskMode::kCustom, DefaultAttention, Params>( + Params params, nv_bfloat16* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_3.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..dd1f6351e5c1fa2e7676b1675fa31103abd29eaa --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_3.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kMultiItemScoring == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 128, 128, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, DefaultAttention, Params>( + Params params, nv_bfloat16* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90.cu new file mode 100644 index 0000000000000000000000000000000000000000..39cabb9c2873d8e3922f6618e7081df62cac1b57 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90.cu @@ -0,0 +1,93 @@ +/* + * Copyright (c) 2024 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include +#include + +#include "pytorch_extension_utils.h" +#include "single_prefill_sm90_config.inc" + +namespace flashinfer { + +template +cudaError_t SinglePrefillWithKVCacheDispatched(Params& params, cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +void single_prefill_with_kv_cache_sm90(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS) { + unsigned int head_dim_qk = q.size(2); + unsigned int head_dim_vo = v.size(2); + unsigned int num_qo_heads = q.size(1); + unsigned int qo_len = q.size(0); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + QKVLayout kv_layout = static_cast(layout); + const c10::cuda::OptionalCUDAGuard device_guard(q.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + const MaskMode mask_mode = static_cast(mask_mode_code); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, USE_SLIDING_WINDOW, + USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + Params params; + params.q_ptr = static_cast(q.data_ptr()); + params.k_ptr = static_cast(k.data_ptr()); + params.v_ptr = static_cast(v.data_ptr()); + params.o_ptr = static_cast(o.data_ptr()); + params.lse_ptr = maybe_lse ? (static_cast(maybe_lse->data_ptr())) : nullptr; + params.q_stride_n = q.stride(0); + params.q_stride_h = q.stride(1); + params.o_stride_n = o.stride(0); + params.o_stride_h = o.stride(1); + if (kv_layout == QKVLayout::kNHD) { + params.k_stride_n = k.stride(0); + params.k_stride_h = k.stride(1); + params.v_stride_n = v.stride(0); + params.v_stride_h = v.stride(1); + } else { + params.k_stride_h = k.stride(0); + params.k_stride_n = k.stride(1); + params.v_stride_h = v.stride(0); + params.v_stride_n = v.stride(1); + } + params.qo_len = q.size(0); + params.kv_len = k.size(0); + params.num_qo_heads = q.size(1); + params.num_kv_heads = k.size(1); + params.causal = mask_mode == MaskMode::kCausal; + params.group_size = params.num_qo_heads / params.num_kv_heads; + params.window_left = window_left; + + ADDITIONAL_PARAMS_SETTER + + cudaError_t status = + SinglePrefillWithKVCacheDispatched(params, + stream); + TORCH_CHECK(status == cudaSuccess, "single_prefill_with_kv_cache_sm90 failed with error: " + + std::string(cudaGetErrorString(status))); + return true; + }); +} diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_config.inc b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..72dfcfdb85915c68467c4a0ffae31520ce70afda --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_config.inc @@ -0,0 +1,70 @@ +#pragma once +#include +#include +#include +#include +#include +#include +#include + + +#define ADDITIONAL_FUNC_PARAMS , double logits_soft_cap, double sm_scale +#define ADDITIONAL_PARAMS_SETTER params.additional_params.logits_soft_cap = logits_soft_cap; \ +params.additional_params.sm_scale = sm_scale; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { using AttentionVariant = DefaultAttention; __VA_ARGS__(); }) + +using namespace flashinfer; + +using DTypeQ = cutlass_dtype_t; +using DTypeKV = cutlass_dtype_t; +using DTypeO = cutlass_dtype_t; +using IdType = cutlass_dtype_t; + +constexpr int HEAD_DIM_QK = 128; +constexpr int HEAD_DIM_VO = 128; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto USE_SLIDING_WINDOW = false; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + // The QKV matrices. + DTypeQ* q_ptr; + DTypeKV* k_ptr; + DTypeKV* v_ptr; + DTypeO* o_ptr; + float* lse_ptr; + + // Additional params + struct AdditionalParams { + double logits_soft_cap; +double sm_scale; +; + } additional_params; + + int64_t q_stride_n; + int64_t k_stride_n; + int64_t v_stride_n; + int64_t o_stride_n; + int64_t q_stride_h; + int64_t k_stride_h; + int64_t v_stride_h; + int64_t o_stride_h; + + int qo_len; + int kv_len; + int head_dim; + int num_qo_heads; + int num_kv_heads; + int group_size; + int window_left; + + bool causal; +}; + +#include \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_jit_pybind.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..2d6dd6076239658b22af333487f1f0fb49a86b1b --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_jit_pybind.cu @@ -0,0 +1,27 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "pytorch_extension_utils.h" +#include "single_prefill_sm90_config.inc" + +void single_prefill_with_kv_cache_sm90(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Single-request prefill attention with KV-Cache operator + m.def("run", single_prefill_with_kv_cache_sm90); +} diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_0.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..1f0930fe51f67c4b1a805a7301222d648998bb0c --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_0.cu @@ -0,0 +1,12 @@ +#include +#include "single_prefill_sm90_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SinglePrefillWithKVCacheDispatched + <128, 128, MaskMode::kNone, /*USE_SLIDING_WINDOW=*/false, DefaultAttention, Params>( + Params& params, cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_1.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..491c13341894b48e93c5e2da9d2a66160ac7bdfe --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_1.cu @@ -0,0 +1,12 @@ +#include +#include "single_prefill_sm90_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SinglePrefillWithKVCacheDispatched + <128, 128, MaskMode::kCausal, /*USE_SLIDING_WINDOW=*/false, DefaultAttention, Params>( + Params& params, cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_2.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..542f3f390d1514cbee2583533d9c269140d340ed --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_2.cu @@ -0,0 +1,12 @@ +#include +#include "single_prefill_sm90_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SinglePrefillWithKVCacheDispatched + <128, 128, MaskMode::kCustom, /*USE_SLIDING_WINDOW=*/false, DefaultAttention, Params>( + Params& params, cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_3.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..981f41acd163d7ef539b7750bd1ec43855e68db4 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_3.cu @@ -0,0 +1,12 @@ +#include +#include "single_prefill_sm90_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SinglePrefillWithKVCacheDispatched + <128, 128, MaskMode::kMultiItemScoring, /*USE_SLIDING_WINDOW=*/false, DefaultAttention, Params>( + Params& params, cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90.cu new file mode 100644 index 0000000000000000000000000000000000000000..39cabb9c2873d8e3922f6618e7081df62cac1b57 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90.cu @@ -0,0 +1,93 @@ +/* + * Copyright (c) 2024 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include +#include + +#include "pytorch_extension_utils.h" +#include "single_prefill_sm90_config.inc" + +namespace flashinfer { + +template +cudaError_t SinglePrefillWithKVCacheDispatched(Params& params, cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +void single_prefill_with_kv_cache_sm90(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS) { + unsigned int head_dim_qk = q.size(2); + unsigned int head_dim_vo = v.size(2); + unsigned int num_qo_heads = q.size(1); + unsigned int qo_len = q.size(0); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + QKVLayout kv_layout = static_cast(layout); + const c10::cuda::OptionalCUDAGuard device_guard(q.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + const MaskMode mask_mode = static_cast(mask_mode_code); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, USE_SLIDING_WINDOW, + USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + Params params; + params.q_ptr = static_cast(q.data_ptr()); + params.k_ptr = static_cast(k.data_ptr()); + params.v_ptr = static_cast(v.data_ptr()); + params.o_ptr = static_cast(o.data_ptr()); + params.lse_ptr = maybe_lse ? (static_cast(maybe_lse->data_ptr())) : nullptr; + params.q_stride_n = q.stride(0); + params.q_stride_h = q.stride(1); + params.o_stride_n = o.stride(0); + params.o_stride_h = o.stride(1); + if (kv_layout == QKVLayout::kNHD) { + params.k_stride_n = k.stride(0); + params.k_stride_h = k.stride(1); + params.v_stride_n = v.stride(0); + params.v_stride_h = v.stride(1); + } else { + params.k_stride_h = k.stride(0); + params.k_stride_n = k.stride(1); + params.v_stride_h = v.stride(0); + params.v_stride_n = v.stride(1); + } + params.qo_len = q.size(0); + params.kv_len = k.size(0); + params.num_qo_heads = q.size(1); + params.num_kv_heads = k.size(1); + params.causal = mask_mode == MaskMode::kCausal; + params.group_size = params.num_qo_heads / params.num_kv_heads; + params.window_left = window_left; + + ADDITIONAL_PARAMS_SETTER + + cudaError_t status = + SinglePrefillWithKVCacheDispatched(params, + stream); + TORCH_CHECK(status == cudaSuccess, "single_prefill_with_kv_cache_sm90 failed with error: " + + std::string(cudaGetErrorString(status))); + return true; + }); +} diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_config.inc b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..1d898b1fbc3e1343ddd0e04ca59da18d303d07b5 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_config.inc @@ -0,0 +1,70 @@ +#pragma once +#include +#include +#include +#include +#include +#include +#include + + +#define ADDITIONAL_FUNC_PARAMS , double logits_soft_cap, double sm_scale +#define ADDITIONAL_PARAMS_SETTER params.additional_params.logits_soft_cap = logits_soft_cap; \ +params.additional_params.sm_scale = sm_scale; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { using AttentionVariant = DefaultAttention; __VA_ARGS__(); }) + +using namespace flashinfer; + +using DTypeQ = cutlass_dtype_t; +using DTypeKV = cutlass_dtype_t; +using DTypeO = cutlass_dtype_t; +using IdType = cutlass_dtype_t; + +constexpr int HEAD_DIM_QK = 192; +constexpr int HEAD_DIM_VO = 128; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto USE_SLIDING_WINDOW = false; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + // The QKV matrices. + DTypeQ* q_ptr; + DTypeKV* k_ptr; + DTypeKV* v_ptr; + DTypeO* o_ptr; + float* lse_ptr; + + // Additional params + struct AdditionalParams { + double logits_soft_cap; +double sm_scale; +; + } additional_params; + + int64_t q_stride_n; + int64_t k_stride_n; + int64_t v_stride_n; + int64_t o_stride_n; + int64_t q_stride_h; + int64_t k_stride_h; + int64_t v_stride_h; + int64_t o_stride_h; + + int qo_len; + int kv_len; + int head_dim; + int num_qo_heads; + int num_kv_heads; + int group_size; + int window_left; + + bool causal; +}; + +#include \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_jit_pybind.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..2d6dd6076239658b22af333487f1f0fb49a86b1b --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_jit_pybind.cu @@ -0,0 +1,27 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "pytorch_extension_utils.h" +#include "single_prefill_sm90_config.inc" + +void single_prefill_with_kv_cache_sm90(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Single-request prefill attention with KV-Cache operator + m.def("run", single_prefill_with_kv_cache_sm90); +} diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_0.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..1a1584eebbf276192085ee7fda13999d2e07580f --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_0.cu @@ -0,0 +1,12 @@ +#include +#include "single_prefill_sm90_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SinglePrefillWithKVCacheDispatched + <192, 128, MaskMode::kNone, /*USE_SLIDING_WINDOW=*/false, DefaultAttention, Params>( + Params& params, cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_1.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..a0b314ab23075044a3f5367515796174e6a3391e --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_1.cu @@ -0,0 +1,12 @@ +#include +#include "single_prefill_sm90_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SinglePrefillWithKVCacheDispatched + <192, 128, MaskMode::kCausal, /*USE_SLIDING_WINDOW=*/false, DefaultAttention, Params>( + Params& params, cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_2.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..020022cd19ca3aaa434424349beede0edf06bbac --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_2.cu @@ -0,0 +1,12 @@ +#include +#include "single_prefill_sm90_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SinglePrefillWithKVCacheDispatched + <192, 128, MaskMode::kCustom, /*USE_SLIDING_WINDOW=*/false, DefaultAttention, Params>( + Params& params, cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_3.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..7eb0d82020eec67bc89f0cabf40c492aa055e3f5 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_3.cu @@ -0,0 +1,12 @@ +#include +#include "single_prefill_sm90_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SinglePrefillWithKVCacheDispatched + <192, 128, MaskMode::kMultiItemScoring, /*USE_SLIDING_WINDOW=*/false, DefaultAttention, Params>( + Params& params, cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill.cu new file mode 100644 index 0000000000000000000000000000000000000000..8d77f8420680e70a7929d8b16bd95f5587a01752 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill.cu @@ -0,0 +1,114 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include + +#include "flashinfer/fastdiv.cuh" +#include "pytorch_extension_utils.h" +#include "single_prefill_config.inc" + +namespace flashinfer { + +template +cudaError_t SinglePrefillWithKVCacheDispatched(Params params, typename Params::DTypeO* tmp, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +void single_prefill_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS) { + auto device = q.device(); + unsigned int head_dim_qk = q.size(2); + unsigned int kv_len, qo_len, num_kv_heads, num_qo_heads; + QKVLayout kv_layout = static_cast(layout); + qo_len = q.size(0); + num_qo_heads = q.size(1); + uint32_t q_stride_n = q.stride(0), q_stride_h = q.stride(1), k_stride_n, k_stride_h, v_stride_n, + v_stride_h; + if (kv_layout == QKVLayout::kNHD) { + kv_len = k.size(0); + num_kv_heads = k.size(1); + k_stride_n = k.stride(0); + k_stride_h = k.stride(1); + v_stride_n = v.stride(0); + v_stride_h = v.stride(1); + } else { + kv_len = k.size(1); + num_kv_heads = k.size(0); + k_stride_h = k.stride(0); + k_stride_n = k.stride(1); + v_stride_h = v.stride(0); + v_stride_n = v.stride(1); + } + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == qo_len, lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == num_qo_heads, lse.size(1), q.size(1)); + } + + const MaskMode mask_mode = static_cast(mask_mode_code); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(device); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, Params, + [&] { + Params params; + + params.q = static_cast(q.data_ptr()); + params.k = static_cast(k.data_ptr()); + params.v = static_cast(v.data_ptr()); + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.num_qo_heads = num_qo_heads; + params.num_kv_heads = num_kv_heads; + params.group_size = uint_fastdiv(num_qo_heads / num_kv_heads); + params.qo_len = qo_len; + params.kv_len = kv_len; + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.k_stride_n = k_stride_n; + params.k_stride_h = k_stride_h; + params.v_stride_n = v_stride_n; + params.v_stride_h = v_stride_h; + + params.window_left = window_left; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + cudaError_t status = flashinfer::SinglePrefillWithKVCacheDispatched< + HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + /*use_fp16_qk_reduction=*/USE_FP16_QK_REDUCTION, MASK_MODE, AttentionVariant>( + params, static_cast(tmp.data_ptr()), stream); + TORCH_CHECK(status == cudaSuccess, + "SinglePrefillWithKVCache kernel launch failed, error: " + + std::string(cudaGetErrorString(status))); + return true; + }); +} diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_config.inc b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..1bee06352bf5d95d64e02de6add67668a17fff1c --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_config.inc @@ -0,0 +1,83 @@ +#pragma once +#include +#include +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_custom_mask, std::optional maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta +#define ADDITIONAL_PARAMS_SETTER params.maybe_custom_mask = maybe_custom_mask ? static_cast(maybe_custom_mask->data_ptr()): nullptr; \ +params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; + + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, Params, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { \ + constexpr bool use_custom_mask = MASK_MODE == MaskMode::kCustom; \ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ + }) + + +using namespace flashinfer; + +using DTypeQ = nv_bfloat16; +using DTypeKV = nv_bfloat16; +using DTypeO = nv_bfloat16; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 256; +constexpr int HEAD_DIM_VO = 256; +constexpr bool USE_FP16_QK_REDUCTION = false; +constexpr auto USE_LOGITS_SOFT_CAP = true; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = true; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = int32_t; + DTypeQ* q; + DTypeKV* k; + DTypeKV* v; + DTypeO* o; + float* lse; + uint_fastdiv group_size; + + uint8_t* maybe_custom_mask; +float* maybe_alibi_slopes; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; + + + uint32_t qo_len; + uint32_t kv_len; + uint32_t num_qo_heads; + uint32_t num_kv_heads; + uint32_t q_stride_n; + uint32_t q_stride_h; + uint32_t k_stride_n; + uint32_t k_stride_h; + uint32_t v_stride_n; + uint32_t v_stride_h; + uint32_t head_dim; + int32_t window_left; + + bool partition_kv; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { + return qo_len; + } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return kv_len; + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_jit_pybind.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..aaccbc045f9fd8566654a106379c99a053489fc7 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_jit_pybind.cu @@ -0,0 +1,27 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "pytorch_extension_utils.h" +#include "single_prefill_config.inc" + +void single_prefill_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Single-request prefill attention with KV-Cache operator + m.def("run", single_prefill_with_kv_cache); +} diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_kernel_mask_0.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..4de45b7f5562f0e0247f7aa1259b039aff99cf44 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_kernel_mask_0.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kNone == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 256, 256, PosEncodingMode::kNone, false, MaskMode::kNone, DefaultAttention, Params>( + Params params, nv_bfloat16* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_kernel_mask_1.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..a8b8b003aeaaff6894cc408701534a38a5a88bec --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_kernel_mask_1.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCausal == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 256, 256, PosEncodingMode::kNone, false, MaskMode::kCausal, DefaultAttention, Params>( + Params params, nv_bfloat16* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_kernel_mask_2.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..64074091bebe3bc9137d6c87d62534531fddcd92 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_kernel_mask_2.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCustom == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 256, 256, PosEncodingMode::kNone, false, MaskMode::kCustom, DefaultAttention, Params>( + Params params, nv_bfloat16* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_kernel_mask_3.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..385a4bcb6d68da17ad291f26d795c705dddbdbc8 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_kernel_mask_3.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kMultiItemScoring == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 256, 256, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, DefaultAttention, Params>( + Params params, nv_bfloat16* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90.cu new file mode 100644 index 0000000000000000000000000000000000000000..39cabb9c2873d8e3922f6618e7081df62cac1b57 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90.cu @@ -0,0 +1,93 @@ +/* + * Copyright (c) 2024 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include +#include + +#include "pytorch_extension_utils.h" +#include "single_prefill_sm90_config.inc" + +namespace flashinfer { + +template +cudaError_t SinglePrefillWithKVCacheDispatched(Params& params, cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +void single_prefill_with_kv_cache_sm90(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS) { + unsigned int head_dim_qk = q.size(2); + unsigned int head_dim_vo = v.size(2); + unsigned int num_qo_heads = q.size(1); + unsigned int qo_len = q.size(0); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + QKVLayout kv_layout = static_cast(layout); + const c10::cuda::OptionalCUDAGuard device_guard(q.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + const MaskMode mask_mode = static_cast(mask_mode_code); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, USE_SLIDING_WINDOW, + USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + Params params; + params.q_ptr = static_cast(q.data_ptr()); + params.k_ptr = static_cast(k.data_ptr()); + params.v_ptr = static_cast(v.data_ptr()); + params.o_ptr = static_cast(o.data_ptr()); + params.lse_ptr = maybe_lse ? (static_cast(maybe_lse->data_ptr())) : nullptr; + params.q_stride_n = q.stride(0); + params.q_stride_h = q.stride(1); + params.o_stride_n = o.stride(0); + params.o_stride_h = o.stride(1); + if (kv_layout == QKVLayout::kNHD) { + params.k_stride_n = k.stride(0); + params.k_stride_h = k.stride(1); + params.v_stride_n = v.stride(0); + params.v_stride_h = v.stride(1); + } else { + params.k_stride_h = k.stride(0); + params.k_stride_n = k.stride(1); + params.v_stride_h = v.stride(0); + params.v_stride_n = v.stride(1); + } + params.qo_len = q.size(0); + params.kv_len = k.size(0); + params.num_qo_heads = q.size(1); + params.num_kv_heads = k.size(1); + params.causal = mask_mode == MaskMode::kCausal; + params.group_size = params.num_qo_heads / params.num_kv_heads; + params.window_left = window_left; + + ADDITIONAL_PARAMS_SETTER + + cudaError_t status = + SinglePrefillWithKVCacheDispatched(params, + stream); + TORCH_CHECK(status == cudaSuccess, "single_prefill_with_kv_cache_sm90 failed with error: " + + std::string(cudaGetErrorString(status))); + return true; + }); +} diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_config.inc b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..238b5f50a5f3aab4a2d162f87790633b018b65e1 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_config.inc @@ -0,0 +1,70 @@ +#pragma once +#include +#include +#include +#include +#include +#include +#include + + +#define ADDITIONAL_FUNC_PARAMS , double logits_soft_cap, double sm_scale +#define ADDITIONAL_PARAMS_SETTER params.additional_params.logits_soft_cap = logits_soft_cap; \ +params.additional_params.sm_scale = sm_scale; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { using AttentionVariant = DefaultAttention; __VA_ARGS__(); }) + +using namespace flashinfer; + +using DTypeQ = cutlass_dtype_t; +using DTypeKV = cutlass_dtype_t; +using DTypeO = cutlass_dtype_t; +using IdType = cutlass_dtype_t; + +constexpr int HEAD_DIM_QK = 256; +constexpr int HEAD_DIM_VO = 256; +constexpr auto USE_LOGITS_SOFT_CAP = true; +constexpr auto USE_SLIDING_WINDOW = true; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + // The QKV matrices. + DTypeQ* q_ptr; + DTypeKV* k_ptr; + DTypeKV* v_ptr; + DTypeO* o_ptr; + float* lse_ptr; + + // Additional params + struct AdditionalParams { + double logits_soft_cap; +double sm_scale; +; + } additional_params; + + int64_t q_stride_n; + int64_t k_stride_n; + int64_t v_stride_n; + int64_t o_stride_n; + int64_t q_stride_h; + int64_t k_stride_h; + int64_t v_stride_h; + int64_t o_stride_h; + + int qo_len; + int kv_len; + int head_dim; + int num_qo_heads; + int num_kv_heads; + int group_size; + int window_left; + + bool causal; +}; + +#include \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_jit_pybind.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..2d6dd6076239658b22af333487f1f0fb49a86b1b --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_jit_pybind.cu @@ -0,0 +1,27 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "pytorch_extension_utils.h" +#include "single_prefill_sm90_config.inc" + +void single_prefill_with_kv_cache_sm90(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Single-request prefill attention with KV-Cache operator + m.def("run", single_prefill_with_kv_cache_sm90); +} diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_kernel_mask_0.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..5b881dbe8c829acac6560accc4ea0a629730ae05 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_kernel_mask_0.cu @@ -0,0 +1,12 @@ +#include +#include "single_prefill_sm90_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SinglePrefillWithKVCacheDispatched + <256, 256, MaskMode::kNone, /*USE_SLIDING_WINDOW=*/true, DefaultAttention, Params>( + Params& params, cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_kernel_mask_1.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..e081f468769e66fd45aef98e65f0c93f5f667745 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_kernel_mask_1.cu @@ -0,0 +1,12 @@ +#include +#include "single_prefill_sm90_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SinglePrefillWithKVCacheDispatched + <256, 256, MaskMode::kCausal, /*USE_SLIDING_WINDOW=*/true, DefaultAttention, Params>( + Params& params, cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_kernel_mask_2.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..ac76c050bef07556fb25a86da5c1c6148d856db5 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_kernel_mask_2.cu @@ -0,0 +1,12 @@ +#include +#include "single_prefill_sm90_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SinglePrefillWithKVCacheDispatched + <256, 256, MaskMode::kCustom, /*USE_SLIDING_WINDOW=*/true, DefaultAttention, Params>( + Params& params, cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_kernel_mask_3.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..6b5e45f1881bea91f4154b5d0968cc5a5750176d --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_kernel_mask_3.cu @@ -0,0 +1,12 @@ +#include +#include "single_prefill_sm90_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SinglePrefillWithKVCacheDispatched + <256, 256, MaskMode::kMultiItemScoring, /*USE_SLIDING_WINDOW=*/true, DefaultAttention, Params>( + Params& params, cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill.cu new file mode 100644 index 0000000000000000000000000000000000000000..8d77f8420680e70a7929d8b16bd95f5587a01752 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill.cu @@ -0,0 +1,114 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include + +#include "flashinfer/fastdiv.cuh" +#include "pytorch_extension_utils.h" +#include "single_prefill_config.inc" + +namespace flashinfer { + +template +cudaError_t SinglePrefillWithKVCacheDispatched(Params params, typename Params::DTypeO* tmp, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +void single_prefill_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS) { + auto device = q.device(); + unsigned int head_dim_qk = q.size(2); + unsigned int kv_len, qo_len, num_kv_heads, num_qo_heads; + QKVLayout kv_layout = static_cast(layout); + qo_len = q.size(0); + num_qo_heads = q.size(1); + uint32_t q_stride_n = q.stride(0), q_stride_h = q.stride(1), k_stride_n, k_stride_h, v_stride_n, + v_stride_h; + if (kv_layout == QKVLayout::kNHD) { + kv_len = k.size(0); + num_kv_heads = k.size(1); + k_stride_n = k.stride(0); + k_stride_h = k.stride(1); + v_stride_n = v.stride(0); + v_stride_h = v.stride(1); + } else { + kv_len = k.size(1); + num_kv_heads = k.size(0); + k_stride_h = k.stride(0); + k_stride_n = k.stride(1); + v_stride_h = v.stride(0); + v_stride_n = v.stride(1); + } + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == qo_len, lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == num_qo_heads, lse.size(1), q.size(1)); + } + + const MaskMode mask_mode = static_cast(mask_mode_code); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(device); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, Params, + [&] { + Params params; + + params.q = static_cast(q.data_ptr()); + params.k = static_cast(k.data_ptr()); + params.v = static_cast(v.data_ptr()); + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.num_qo_heads = num_qo_heads; + params.num_kv_heads = num_kv_heads; + params.group_size = uint_fastdiv(num_qo_heads / num_kv_heads); + params.qo_len = qo_len; + params.kv_len = kv_len; + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.k_stride_n = k_stride_n; + params.k_stride_h = k_stride_h; + params.v_stride_n = v_stride_n; + params.v_stride_h = v_stride_h; + + params.window_left = window_left; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + cudaError_t status = flashinfer::SinglePrefillWithKVCacheDispatched< + HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + /*use_fp16_qk_reduction=*/USE_FP16_QK_REDUCTION, MASK_MODE, AttentionVariant>( + params, static_cast(tmp.data_ptr()), stream); + TORCH_CHECK(status == cudaSuccess, + "SinglePrefillWithKVCache kernel launch failed, error: " + + std::string(cudaGetErrorString(status))); + return true; + }); +} diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_config.inc b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..ca64519911a34a7d5e9a412f4dde847e8cc18e54 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_config.inc @@ -0,0 +1,83 @@ +#pragma once +#include +#include +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_custom_mask, std::optional maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta +#define ADDITIONAL_PARAMS_SETTER params.maybe_custom_mask = maybe_custom_mask ? static_cast(maybe_custom_mask->data_ptr()): nullptr; \ +params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; + + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, Params, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { \ + constexpr bool use_custom_mask = MASK_MODE == MaskMode::kCustom; \ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ + }) + + +using namespace flashinfer; + +using DTypeQ = nv_bfloat16; +using DTypeKV = nv_bfloat16; +using DTypeO = nv_bfloat16; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 64; +constexpr int HEAD_DIM_VO = 64; +constexpr bool USE_FP16_QK_REDUCTION = false; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = false; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = int32_t; + DTypeQ* q; + DTypeKV* k; + DTypeKV* v; + DTypeO* o; + float* lse; + uint_fastdiv group_size; + + uint8_t* maybe_custom_mask; +float* maybe_alibi_slopes; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; + + + uint32_t qo_len; + uint32_t kv_len; + uint32_t num_qo_heads; + uint32_t num_kv_heads; + uint32_t q_stride_n; + uint32_t q_stride_h; + uint32_t k_stride_n; + uint32_t k_stride_h; + uint32_t v_stride_n; + uint32_t v_stride_h; + uint32_t head_dim; + int32_t window_left; + + bool partition_kv; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { + return qo_len; + } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return kv_len; + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_jit_pybind.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..aaccbc045f9fd8566654a106379c99a053489fc7 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_jit_pybind.cu @@ -0,0 +1,27 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "pytorch_extension_utils.h" +#include "single_prefill_config.inc" + +void single_prefill_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Single-request prefill attention with KV-Cache operator + m.def("run", single_prefill_with_kv_cache); +} diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_0.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..80c431f71ae682b5a984273dc3f21a148695295e --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_0.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kNone == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, DefaultAttention, Params>( + Params params, nv_bfloat16* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_1.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..fa50e323028ecfadb643d4a44d4b3400f793add8 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_1.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCausal == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, DefaultAttention, Params>( + Params params, nv_bfloat16* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_2.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..152acf34b610ae996d4d3411fc8d6ddbb114bab0 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_2.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCustom == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, DefaultAttention, Params>( + Params params, nv_bfloat16* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_3.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..5f6b6f6745d4b538063664565da078390c14e7ad --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_3.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kMultiItemScoring == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, DefaultAttention, Params>( + Params params, nv_bfloat16* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill.cu new file mode 100644 index 0000000000000000000000000000000000000000..8d77f8420680e70a7929d8b16bd95f5587a01752 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill.cu @@ -0,0 +1,114 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include + +#include "flashinfer/fastdiv.cuh" +#include "pytorch_extension_utils.h" +#include "single_prefill_config.inc" + +namespace flashinfer { + +template +cudaError_t SinglePrefillWithKVCacheDispatched(Params params, typename Params::DTypeO* tmp, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +void single_prefill_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS) { + auto device = q.device(); + unsigned int head_dim_qk = q.size(2); + unsigned int kv_len, qo_len, num_kv_heads, num_qo_heads; + QKVLayout kv_layout = static_cast(layout); + qo_len = q.size(0); + num_qo_heads = q.size(1); + uint32_t q_stride_n = q.stride(0), q_stride_h = q.stride(1), k_stride_n, k_stride_h, v_stride_n, + v_stride_h; + if (kv_layout == QKVLayout::kNHD) { + kv_len = k.size(0); + num_kv_heads = k.size(1); + k_stride_n = k.stride(0); + k_stride_h = k.stride(1); + v_stride_n = v.stride(0); + v_stride_h = v.stride(1); + } else { + kv_len = k.size(1); + num_kv_heads = k.size(0); + k_stride_h = k.stride(0); + k_stride_n = k.stride(1); + v_stride_h = v.stride(0); + v_stride_n = v.stride(1); + } + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == qo_len, lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == num_qo_heads, lse.size(1), q.size(1)); + } + + const MaskMode mask_mode = static_cast(mask_mode_code); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(device); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, Params, + [&] { + Params params; + + params.q = static_cast(q.data_ptr()); + params.k = static_cast(k.data_ptr()); + params.v = static_cast(v.data_ptr()); + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.num_qo_heads = num_qo_heads; + params.num_kv_heads = num_kv_heads; + params.group_size = uint_fastdiv(num_qo_heads / num_kv_heads); + params.qo_len = qo_len; + params.kv_len = kv_len; + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.k_stride_n = k_stride_n; + params.k_stride_h = k_stride_h; + params.v_stride_n = v_stride_n; + params.v_stride_h = v_stride_h; + + params.window_left = window_left; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + cudaError_t status = flashinfer::SinglePrefillWithKVCacheDispatched< + HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + /*use_fp16_qk_reduction=*/USE_FP16_QK_REDUCTION, MASK_MODE, AttentionVariant>( + params, static_cast(tmp.data_ptr()), stream); + TORCH_CHECK(status == cudaSuccess, + "SinglePrefillWithKVCache kernel launch failed, error: " + + std::string(cudaGetErrorString(status))); + return true; + }); +} diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_config.inc b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..4482275be962c1e9de23d6207dc1e20a139556f7 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_config.inc @@ -0,0 +1,83 @@ +#pragma once +#include +#include +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_custom_mask, std::optional maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta +#define ADDITIONAL_PARAMS_SETTER params.maybe_custom_mask = maybe_custom_mask ? static_cast(maybe_custom_mask->data_ptr()): nullptr; \ +params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; + + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, Params, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { \ + constexpr bool use_custom_mask = MASK_MODE == MaskMode::kCustom; \ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ + }) + + +using namespace flashinfer; + +using DTypeQ = nv_bfloat16; +using DTypeKV = nv_bfloat16; +using DTypeO = nv_bfloat16; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 64; +constexpr int HEAD_DIM_VO = 64; +constexpr bool USE_FP16_QK_REDUCTION = false; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = true; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = int32_t; + DTypeQ* q; + DTypeKV* k; + DTypeKV* v; + DTypeO* o; + float* lse; + uint_fastdiv group_size; + + uint8_t* maybe_custom_mask; +float* maybe_alibi_slopes; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; + + + uint32_t qo_len; + uint32_t kv_len; + uint32_t num_qo_heads; + uint32_t num_kv_heads; + uint32_t q_stride_n; + uint32_t q_stride_h; + uint32_t k_stride_n; + uint32_t k_stride_h; + uint32_t v_stride_n; + uint32_t v_stride_h; + uint32_t head_dim; + int32_t window_left; + + bool partition_kv; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { + return qo_len; + } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return kv_len; + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_jit_pybind.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..aaccbc045f9fd8566654a106379c99a053489fc7 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_jit_pybind.cu @@ -0,0 +1,27 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "pytorch_extension_utils.h" +#include "single_prefill_config.inc" + +void single_prefill_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Single-request prefill attention with KV-Cache operator + m.def("run", single_prefill_with_kv_cache); +} diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_0.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..ff59a581dc69eac31ab74b980d4f7ce073944a92 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_0.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kNone == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, DefaultAttention, Params>( + Params params, nv_bfloat16* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_1.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..fb95a10fd40a9e60fb9d061a43de09a4f672322f --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_1.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCausal == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, DefaultAttention, Params>( + Params params, nv_bfloat16* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_2.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..cd6110a5a31b900002d0d6be9bc426a6291c3b72 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_2.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCustom == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, DefaultAttention, Params>( + Params params, nv_bfloat16* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_3.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..db74db3b3364a6185542af8077ba92cc368418eb --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_bf16_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_3.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kMultiItemScoring == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, DefaultAttention, Params>( + Params params, nv_bfloat16* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill.cu new file mode 100644 index 0000000000000000000000000000000000000000..8d77f8420680e70a7929d8b16bd95f5587a01752 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill.cu @@ -0,0 +1,114 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include + +#include "flashinfer/fastdiv.cuh" +#include "pytorch_extension_utils.h" +#include "single_prefill_config.inc" + +namespace flashinfer { + +template +cudaError_t SinglePrefillWithKVCacheDispatched(Params params, typename Params::DTypeO* tmp, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +void single_prefill_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS) { + auto device = q.device(); + unsigned int head_dim_qk = q.size(2); + unsigned int kv_len, qo_len, num_kv_heads, num_qo_heads; + QKVLayout kv_layout = static_cast(layout); + qo_len = q.size(0); + num_qo_heads = q.size(1); + uint32_t q_stride_n = q.stride(0), q_stride_h = q.stride(1), k_stride_n, k_stride_h, v_stride_n, + v_stride_h; + if (kv_layout == QKVLayout::kNHD) { + kv_len = k.size(0); + num_kv_heads = k.size(1); + k_stride_n = k.stride(0); + k_stride_h = k.stride(1); + v_stride_n = v.stride(0); + v_stride_h = v.stride(1); + } else { + kv_len = k.size(1); + num_kv_heads = k.size(0); + k_stride_h = k.stride(0); + k_stride_n = k.stride(1); + v_stride_h = v.stride(0); + v_stride_n = v.stride(1); + } + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == qo_len, lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == num_qo_heads, lse.size(1), q.size(1)); + } + + const MaskMode mask_mode = static_cast(mask_mode_code); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(device); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, Params, + [&] { + Params params; + + params.q = static_cast(q.data_ptr()); + params.k = static_cast(k.data_ptr()); + params.v = static_cast(v.data_ptr()); + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.num_qo_heads = num_qo_heads; + params.num_kv_heads = num_kv_heads; + params.group_size = uint_fastdiv(num_qo_heads / num_kv_heads); + params.qo_len = qo_len; + params.kv_len = kv_len; + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.k_stride_n = k_stride_n; + params.k_stride_h = k_stride_h; + params.v_stride_n = v_stride_n; + params.v_stride_h = v_stride_h; + + params.window_left = window_left; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + cudaError_t status = flashinfer::SinglePrefillWithKVCacheDispatched< + HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + /*use_fp16_qk_reduction=*/USE_FP16_QK_REDUCTION, MASK_MODE, AttentionVariant>( + params, static_cast(tmp.data_ptr()), stream); + TORCH_CHECK(status == cudaSuccess, + "SinglePrefillWithKVCache kernel launch failed, error: " + + std::string(cudaGetErrorString(status))); + return true; + }); +} diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_config.inc b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..a7c97295596fb79ff4f3e92c52da298d4bdd936d --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_config.inc @@ -0,0 +1,83 @@ +#pragma once +#include +#include +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_custom_mask, std::optional maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta +#define ADDITIONAL_PARAMS_SETTER params.maybe_custom_mask = maybe_custom_mask ? static_cast(maybe_custom_mask->data_ptr()): nullptr; \ +params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; + + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, Params, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { \ + constexpr bool use_custom_mask = MASK_MODE == MaskMode::kCustom; \ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ + }) + + +using namespace flashinfer; + +using DTypeQ = nv_bfloat16; +using DTypeKV = __nv_fp8_e4m3; +using DTypeO = nv_bfloat16; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 128; +constexpr int HEAD_DIM_VO = 128; +constexpr bool USE_FP16_QK_REDUCTION = false; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = false; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = int32_t; + DTypeQ* q; + DTypeKV* k; + DTypeKV* v; + DTypeO* o; + float* lse; + uint_fastdiv group_size; + + uint8_t* maybe_custom_mask; +float* maybe_alibi_slopes; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; + + + uint32_t qo_len; + uint32_t kv_len; + uint32_t num_qo_heads; + uint32_t num_kv_heads; + uint32_t q_stride_n; + uint32_t q_stride_h; + uint32_t k_stride_n; + uint32_t k_stride_h; + uint32_t v_stride_n; + uint32_t v_stride_h; + uint32_t head_dim; + int32_t window_left; + + bool partition_kv; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { + return qo_len; + } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return kv_len; + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_jit_pybind.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..aaccbc045f9fd8566654a106379c99a053489fc7 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_jit_pybind.cu @@ -0,0 +1,27 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "pytorch_extension_utils.h" +#include "single_prefill_config.inc" + +void single_prefill_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Single-request prefill attention with KV-Cache operator + m.def("run", single_prefill_with_kv_cache); +} diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_0.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..96c9a8e7cf1247b5c741af94c480cf5068a2f97a --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_0.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kNone == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 128, 128, PosEncodingMode::kNone, false, MaskMode::kNone, DefaultAttention, Params>( + Params params, nv_bfloat16* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_1.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..8c4bdb83a9b4e27744f1ca7b8a3b32a72f1bb5f5 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_1.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCausal == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 128, 128, PosEncodingMode::kNone, false, MaskMode::kCausal, DefaultAttention, Params>( + Params params, nv_bfloat16* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_2.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..8d038b88e1988b2313593ad20d0f8fcc108ce129 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_2.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCustom == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 128, 128, PosEncodingMode::kNone, false, MaskMode::kCustom, DefaultAttention, Params>( + Params params, nv_bfloat16* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_3.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..dd1f6351e5c1fa2e7676b1675fa31103abd29eaa --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_3.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kMultiItemScoring == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 128, 128, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, DefaultAttention, Params>( + Params params, nv_bfloat16* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill.cu new file mode 100644 index 0000000000000000000000000000000000000000..8d77f8420680e70a7929d8b16bd95f5587a01752 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill.cu @@ -0,0 +1,114 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include + +#include "flashinfer/fastdiv.cuh" +#include "pytorch_extension_utils.h" +#include "single_prefill_config.inc" + +namespace flashinfer { + +template +cudaError_t SinglePrefillWithKVCacheDispatched(Params params, typename Params::DTypeO* tmp, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +void single_prefill_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS) { + auto device = q.device(); + unsigned int head_dim_qk = q.size(2); + unsigned int kv_len, qo_len, num_kv_heads, num_qo_heads; + QKVLayout kv_layout = static_cast(layout); + qo_len = q.size(0); + num_qo_heads = q.size(1); + uint32_t q_stride_n = q.stride(0), q_stride_h = q.stride(1), k_stride_n, k_stride_h, v_stride_n, + v_stride_h; + if (kv_layout == QKVLayout::kNHD) { + kv_len = k.size(0); + num_kv_heads = k.size(1); + k_stride_n = k.stride(0); + k_stride_h = k.stride(1); + v_stride_n = v.stride(0); + v_stride_h = v.stride(1); + } else { + kv_len = k.size(1); + num_kv_heads = k.size(0); + k_stride_h = k.stride(0); + k_stride_n = k.stride(1); + v_stride_h = v.stride(0); + v_stride_n = v.stride(1); + } + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == qo_len, lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == num_qo_heads, lse.size(1), q.size(1)); + } + + const MaskMode mask_mode = static_cast(mask_mode_code); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(device); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, Params, + [&] { + Params params; + + params.q = static_cast(q.data_ptr()); + params.k = static_cast(k.data_ptr()); + params.v = static_cast(v.data_ptr()); + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.num_qo_heads = num_qo_heads; + params.num_kv_heads = num_kv_heads; + params.group_size = uint_fastdiv(num_qo_heads / num_kv_heads); + params.qo_len = qo_len; + params.kv_len = kv_len; + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.k_stride_n = k_stride_n; + params.k_stride_h = k_stride_h; + params.v_stride_n = v_stride_n; + params.v_stride_h = v_stride_h; + + params.window_left = window_left; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + cudaError_t status = flashinfer::SinglePrefillWithKVCacheDispatched< + HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + /*use_fp16_qk_reduction=*/USE_FP16_QK_REDUCTION, MASK_MODE, AttentionVariant>( + params, static_cast(tmp.data_ptr()), stream); + TORCH_CHECK(status == cudaSuccess, + "SinglePrefillWithKVCache kernel launch failed, error: " + + std::string(cudaGetErrorString(status))); + return true; + }); +} diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_config.inc b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..bcc001e2a702c198dd81e9717b22f8f57141da22 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_config.inc @@ -0,0 +1,83 @@ +#pragma once +#include +#include +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_custom_mask, std::optional maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta +#define ADDITIONAL_PARAMS_SETTER params.maybe_custom_mask = maybe_custom_mask ? static_cast(maybe_custom_mask->data_ptr()): nullptr; \ +params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; + + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, Params, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { \ + constexpr bool use_custom_mask = MASK_MODE == MaskMode::kCustom; \ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ + }) + + +using namespace flashinfer; + +using DTypeQ = nv_bfloat16; +using DTypeKV = __nv_fp8_e4m3; +using DTypeO = nv_bfloat16; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 256; +constexpr int HEAD_DIM_VO = 256; +constexpr bool USE_FP16_QK_REDUCTION = false; +constexpr auto USE_LOGITS_SOFT_CAP = true; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = true; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = int32_t; + DTypeQ* q; + DTypeKV* k; + DTypeKV* v; + DTypeO* o; + float* lse; + uint_fastdiv group_size; + + uint8_t* maybe_custom_mask; +float* maybe_alibi_slopes; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; + + + uint32_t qo_len; + uint32_t kv_len; + uint32_t num_qo_heads; + uint32_t num_kv_heads; + uint32_t q_stride_n; + uint32_t q_stride_h; + uint32_t k_stride_n; + uint32_t k_stride_h; + uint32_t v_stride_n; + uint32_t v_stride_h; + uint32_t head_dim; + int32_t window_left; + + bool partition_kv; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { + return qo_len; + } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return kv_len; + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_jit_pybind.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..aaccbc045f9fd8566654a106379c99a053489fc7 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_jit_pybind.cu @@ -0,0 +1,27 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "pytorch_extension_utils.h" +#include "single_prefill_config.inc" + +void single_prefill_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Single-request prefill attention with KV-Cache operator + m.def("run", single_prefill_with_kv_cache); +} diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_kernel_mask_0.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..4de45b7f5562f0e0247f7aa1259b039aff99cf44 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_kernel_mask_0.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kNone == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 256, 256, PosEncodingMode::kNone, false, MaskMode::kNone, DefaultAttention, Params>( + Params params, nv_bfloat16* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_kernel_mask_1.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..a8b8b003aeaaff6894cc408701534a38a5a88bec --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_kernel_mask_1.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCausal == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 256, 256, PosEncodingMode::kNone, false, MaskMode::kCausal, DefaultAttention, Params>( + Params params, nv_bfloat16* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_kernel_mask_2.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..64074091bebe3bc9137d6c87d62534531fddcd92 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_kernel_mask_2.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCustom == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 256, 256, PosEncodingMode::kNone, false, MaskMode::kCustom, DefaultAttention, Params>( + Params params, nv_bfloat16* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_kernel_mask_3.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..385a4bcb6d68da17ad291f26d795c705dddbdbc8 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_kernel_mask_3.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kMultiItemScoring == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 256, 256, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, DefaultAttention, Params>( + Params params, nv_bfloat16* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill.cu new file mode 100644 index 0000000000000000000000000000000000000000..8d77f8420680e70a7929d8b16bd95f5587a01752 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill.cu @@ -0,0 +1,114 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include + +#include "flashinfer/fastdiv.cuh" +#include "pytorch_extension_utils.h" +#include "single_prefill_config.inc" + +namespace flashinfer { + +template +cudaError_t SinglePrefillWithKVCacheDispatched(Params params, typename Params::DTypeO* tmp, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +void single_prefill_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS) { + auto device = q.device(); + unsigned int head_dim_qk = q.size(2); + unsigned int kv_len, qo_len, num_kv_heads, num_qo_heads; + QKVLayout kv_layout = static_cast(layout); + qo_len = q.size(0); + num_qo_heads = q.size(1); + uint32_t q_stride_n = q.stride(0), q_stride_h = q.stride(1), k_stride_n, k_stride_h, v_stride_n, + v_stride_h; + if (kv_layout == QKVLayout::kNHD) { + kv_len = k.size(0); + num_kv_heads = k.size(1); + k_stride_n = k.stride(0); + k_stride_h = k.stride(1); + v_stride_n = v.stride(0); + v_stride_h = v.stride(1); + } else { + kv_len = k.size(1); + num_kv_heads = k.size(0); + k_stride_h = k.stride(0); + k_stride_n = k.stride(1); + v_stride_h = v.stride(0); + v_stride_n = v.stride(1); + } + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == qo_len, lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == num_qo_heads, lse.size(1), q.size(1)); + } + + const MaskMode mask_mode = static_cast(mask_mode_code); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(device); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, Params, + [&] { + Params params; + + params.q = static_cast(q.data_ptr()); + params.k = static_cast(k.data_ptr()); + params.v = static_cast(v.data_ptr()); + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.num_qo_heads = num_qo_heads; + params.num_kv_heads = num_kv_heads; + params.group_size = uint_fastdiv(num_qo_heads / num_kv_heads); + params.qo_len = qo_len; + params.kv_len = kv_len; + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.k_stride_n = k_stride_n; + params.k_stride_h = k_stride_h; + params.v_stride_n = v_stride_n; + params.v_stride_h = v_stride_h; + + params.window_left = window_left; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + cudaError_t status = flashinfer::SinglePrefillWithKVCacheDispatched< + HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + /*use_fp16_qk_reduction=*/USE_FP16_QK_REDUCTION, MASK_MODE, AttentionVariant>( + params, static_cast(tmp.data_ptr()), stream); + TORCH_CHECK(status == cudaSuccess, + "SinglePrefillWithKVCache kernel launch failed, error: " + + std::string(cudaGetErrorString(status))); + return true; + }); +} diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_config.inc b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..8df2c5951d3bbf6e40ff7e0d8f7499baee7be2c2 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_config.inc @@ -0,0 +1,83 @@ +#pragma once +#include +#include +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_custom_mask, std::optional maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta +#define ADDITIONAL_PARAMS_SETTER params.maybe_custom_mask = maybe_custom_mask ? static_cast(maybe_custom_mask->data_ptr()): nullptr; \ +params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; + + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, Params, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { \ + constexpr bool use_custom_mask = MASK_MODE == MaskMode::kCustom; \ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ + }) + + +using namespace flashinfer; + +using DTypeQ = nv_bfloat16; +using DTypeKV = __nv_fp8_e4m3; +using DTypeO = nv_bfloat16; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 64; +constexpr int HEAD_DIM_VO = 64; +constexpr bool USE_FP16_QK_REDUCTION = false; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = false; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = int32_t; + DTypeQ* q; + DTypeKV* k; + DTypeKV* v; + DTypeO* o; + float* lse; + uint_fastdiv group_size; + + uint8_t* maybe_custom_mask; +float* maybe_alibi_slopes; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; + + + uint32_t qo_len; + uint32_t kv_len; + uint32_t num_qo_heads; + uint32_t num_kv_heads; + uint32_t q_stride_n; + uint32_t q_stride_h; + uint32_t k_stride_n; + uint32_t k_stride_h; + uint32_t v_stride_n; + uint32_t v_stride_h; + uint32_t head_dim; + int32_t window_left; + + bool partition_kv; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { + return qo_len; + } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return kv_len; + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_jit_pybind.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..aaccbc045f9fd8566654a106379c99a053489fc7 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_jit_pybind.cu @@ -0,0 +1,27 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "pytorch_extension_utils.h" +#include "single_prefill_config.inc" + +void single_prefill_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Single-request prefill attention with KV-Cache operator + m.def("run", single_prefill_with_kv_cache); +} diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_0.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..80c431f71ae682b5a984273dc3f21a148695295e --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_0.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kNone == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, DefaultAttention, Params>( + Params params, nv_bfloat16* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_1.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..fa50e323028ecfadb643d4a44d4b3400f793add8 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_1.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCausal == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, DefaultAttention, Params>( + Params params, nv_bfloat16* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_2.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..152acf34b610ae996d4d3411fc8d6ddbb114bab0 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_2.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCustom == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, DefaultAttention, Params>( + Params params, nv_bfloat16* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_3.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..5f6b6f6745d4b538063664565da078390c14e7ad --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_3.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kMultiItemScoring == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, DefaultAttention, Params>( + Params params, nv_bfloat16* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill.cu new file mode 100644 index 0000000000000000000000000000000000000000..8d77f8420680e70a7929d8b16bd95f5587a01752 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill.cu @@ -0,0 +1,114 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include + +#include "flashinfer/fastdiv.cuh" +#include "pytorch_extension_utils.h" +#include "single_prefill_config.inc" + +namespace flashinfer { + +template +cudaError_t SinglePrefillWithKVCacheDispatched(Params params, typename Params::DTypeO* tmp, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +void single_prefill_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS) { + auto device = q.device(); + unsigned int head_dim_qk = q.size(2); + unsigned int kv_len, qo_len, num_kv_heads, num_qo_heads; + QKVLayout kv_layout = static_cast(layout); + qo_len = q.size(0); + num_qo_heads = q.size(1); + uint32_t q_stride_n = q.stride(0), q_stride_h = q.stride(1), k_stride_n, k_stride_h, v_stride_n, + v_stride_h; + if (kv_layout == QKVLayout::kNHD) { + kv_len = k.size(0); + num_kv_heads = k.size(1); + k_stride_n = k.stride(0); + k_stride_h = k.stride(1); + v_stride_n = v.stride(0); + v_stride_h = v.stride(1); + } else { + kv_len = k.size(1); + num_kv_heads = k.size(0); + k_stride_h = k.stride(0); + k_stride_n = k.stride(1); + v_stride_h = v.stride(0); + v_stride_n = v.stride(1); + } + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == qo_len, lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == num_qo_heads, lse.size(1), q.size(1)); + } + + const MaskMode mask_mode = static_cast(mask_mode_code); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(device); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, Params, + [&] { + Params params; + + params.q = static_cast(q.data_ptr()); + params.k = static_cast(k.data_ptr()); + params.v = static_cast(v.data_ptr()); + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.num_qo_heads = num_qo_heads; + params.num_kv_heads = num_kv_heads; + params.group_size = uint_fastdiv(num_qo_heads / num_kv_heads); + params.qo_len = qo_len; + params.kv_len = kv_len; + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.k_stride_n = k_stride_n; + params.k_stride_h = k_stride_h; + params.v_stride_n = v_stride_n; + params.v_stride_h = v_stride_h; + + params.window_left = window_left; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + cudaError_t status = flashinfer::SinglePrefillWithKVCacheDispatched< + HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + /*use_fp16_qk_reduction=*/USE_FP16_QK_REDUCTION, MASK_MODE, AttentionVariant>( + params, static_cast(tmp.data_ptr()), stream); + TORCH_CHECK(status == cudaSuccess, + "SinglePrefillWithKVCache kernel launch failed, error: " + + std::string(cudaGetErrorString(status))); + return true; + }); +} diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_config.inc b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..33ff6e876395f76f868b768c309be7489f71210f --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_config.inc @@ -0,0 +1,83 @@ +#pragma once +#include +#include +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_custom_mask, std::optional maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta +#define ADDITIONAL_PARAMS_SETTER params.maybe_custom_mask = maybe_custom_mask ? static_cast(maybe_custom_mask->data_ptr()): nullptr; \ +params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; + + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, Params, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { \ + constexpr bool use_custom_mask = MASK_MODE == MaskMode::kCustom; \ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ + }) + + +using namespace flashinfer; + +using DTypeQ = nv_bfloat16; +using DTypeKV = __nv_fp8_e4m3; +using DTypeO = nv_bfloat16; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 64; +constexpr int HEAD_DIM_VO = 64; +constexpr bool USE_FP16_QK_REDUCTION = false; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = true; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = int32_t; + DTypeQ* q; + DTypeKV* k; + DTypeKV* v; + DTypeO* o; + float* lse; + uint_fastdiv group_size; + + uint8_t* maybe_custom_mask; +float* maybe_alibi_slopes; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; + + + uint32_t qo_len; + uint32_t kv_len; + uint32_t num_qo_heads; + uint32_t num_kv_heads; + uint32_t q_stride_n; + uint32_t q_stride_h; + uint32_t k_stride_n; + uint32_t k_stride_h; + uint32_t v_stride_n; + uint32_t v_stride_h; + uint32_t head_dim; + int32_t window_left; + + bool partition_kv; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { + return qo_len; + } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return kv_len; + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_jit_pybind.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..aaccbc045f9fd8566654a106379c99a053489fc7 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_jit_pybind.cu @@ -0,0 +1,27 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "pytorch_extension_utils.h" +#include "single_prefill_config.inc" + +void single_prefill_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Single-request prefill attention with KV-Cache operator + m.def("run", single_prefill_with_kv_cache); +} diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_0.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..ff59a581dc69eac31ab74b980d4f7ce073944a92 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_0.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kNone == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, DefaultAttention, Params>( + Params params, nv_bfloat16* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_1.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..fb95a10fd40a9e60fb9d061a43de09a4f672322f --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_1.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCausal == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, DefaultAttention, Params>( + Params params, nv_bfloat16* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_2.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..cd6110a5a31b900002d0d6be9bc426a6291c3b72 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_2.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCustom == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, DefaultAttention, Params>( + Params params, nv_bfloat16* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_3.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..db74db3b3364a6185542af8077ba92cc368418eb --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_bf16_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_3.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kMultiItemScoring == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, DefaultAttention, Params>( + Params params, nv_bfloat16* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_fp8_sm90.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_fp8_sm90.cu new file mode 100644 index 0000000000000000000000000000000000000000..2e4eb9133649c0b339f2352e2e7ae49475c5c99a --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_fp8_sm90.cu @@ -0,0 +1,101 @@ +/* + * Copyright (c) 2024 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include +#include + +#include "pytorch_extension_utils.h" +#include "single_prefill_sm90_config.inc" + +namespace flashinfer { + +template +cudaError_t SingleFP8PrefillWithKVCacheDispatched(Params& params, cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +void single_prefill_with_kv_cache_sm90(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS) { + unsigned int head_dim_qk = q.size(2); + unsigned int head_dim_vo = v.size(2); + unsigned int num_qo_heads = q.size(1); + unsigned int qo_len = q.size(0); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + QKVLayout kv_layout = static_cast(layout); + const c10::cuda::OptionalCUDAGuard device_guard(q.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + const MaskMode mask_mode = static_cast(mask_mode_code); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, USE_SLIDING_WINDOW, + USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + Params params; + params.q_ptr = static_cast(q.data_ptr()); + params.k_ptr = static_cast(k.data_ptr()); + params.v_ptr = static_cast(v.data_ptr()); + params.o_ptr = static_cast(o.data_ptr()); + params.lse_ptr = maybe_lse ? (static_cast(maybe_lse->data_ptr())) : nullptr; + params.q_stride_n = q.stride(0); + params.q_stride_h = q.stride(1); + params.o_stride_n = o.stride(0); + params.o_stride_h = o.stride(1); + if (kv_layout == QKVLayout::kNHD) { + params.k_stride_n = k.stride(0); + params.k_stride_h = k.stride(1); + params.v_stride_n = v.stride(0); + params.v_stride_h = v.stride(1); + } else { + params.k_stride_h = k.stride(0); + params.k_stride_n = k.stride(1); + params.v_stride_h = v.stride(0); + params.v_stride_n = v.stride(1); + } + params.qo_len = q.size(0); + params.kv_len = k.size(0); + params.num_qo_heads = q.size(1); + params.num_kv_heads = k.size(1); + params.causal = mask_mode == MaskMode::kCausal; + params.group_size = params.num_qo_heads / params.num_kv_heads; + + // Note(Yilong): this should be checked on Python Side + // Only support window_left == 0 for now + params.window_left = window_left; + + // Note(Yilong): all quantization parameters are set in additional_params + ADDITIONAL_PARAMS_SETTER + + // Not support various head_dim for now + static_assert(HEAD_DIM_QK == HEAD_DIM_VO, "head_dim_qk and head_dim_vo should be the same"); + // Currently only support same quantization precision + static_assert(std::is_same_v); + + cudaError_t status = + SingleFP8PrefillWithKVCacheDispatched(params, stream); + TORCH_CHECK(status == cudaSuccess, "single_prefill_with_kv_cache_sm90 failed with error: " + + std::string(cudaGetErrorString(status))); + return true; + }); +} diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_config.inc b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..60ad5b74f6692ca09a48f82a2694e1af00dd5302 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_config.inc @@ -0,0 +1,74 @@ +#pragma once +#include +#include +#include +#include +#include +#include +#include + + +#define ADDITIONAL_FUNC_PARAMS , at::Tensor scale_q, at::Tensor scale_k, at::Tensor scale_v, double sm_scale +#define ADDITIONAL_PARAMS_SETTER params.additional_params.scale_q = static_cast(scale_q.data_ptr()); \ +params.additional_params.scale_k = static_cast(scale_k.data_ptr()); \ +params.additional_params.scale_v = static_cast(scale_v.data_ptr()); \ +params.additional_params.sm_scale = sm_scale; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { using AttentionVariant = DefaultFP8Attention; __VA_ARGS__(); }) + +using namespace flashinfer; + +using DTypeQ = cutlass_dtype_t<__nv_fp8_e4m3>; +using DTypeKV = cutlass_dtype_t<__nv_fp8_e4m3>; +using DTypeO = cutlass_dtype_t; +using IdType = cutlass_dtype_t; + +constexpr int HEAD_DIM_QK = 128; +constexpr int HEAD_DIM_VO = 128; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto USE_SLIDING_WINDOW = false; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + // The QKV matrices. + DTypeQ* q_ptr; + DTypeKV* k_ptr; + DTypeKV* v_ptr; + DTypeO* o_ptr; + float* lse_ptr; + + // Additional params + struct AdditionalParams { + float* scale_q; +float* scale_k; +float* scale_v; +double sm_scale; +; + } additional_params; + + int64_t q_stride_n; + int64_t k_stride_n; + int64_t v_stride_n; + int64_t o_stride_n; + int64_t q_stride_h; + int64_t k_stride_h; + int64_t v_stride_h; + int64_t o_stride_h; + + int qo_len; + int kv_len; + int head_dim; + int num_qo_heads; + int num_kv_heads; + int group_size; + int window_left; + + bool causal; +}; + +#include \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_jit_pybind.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..2d6dd6076239658b22af333487f1f0fb49a86b1b --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_jit_pybind.cu @@ -0,0 +1,27 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "pytorch_extension_utils.h" +#include "single_prefill_sm90_config.inc" + +void single_prefill_with_kv_cache_sm90(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Single-request prefill attention with KV-Cache operator + m.def("run", single_prefill_with_kv_cache_sm90); +} diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_0.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..cea7f733ac8b56b095898da8c22d733aa53c7a3b --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_0.cu @@ -0,0 +1,11 @@ +#include +#include "single_prefill_sm90_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SingleFP8PrefillWithKVCacheDispatched + <128, MaskMode::kNone, /*USE_SLIDING_WINDOW=*/false, DefaultFP8Attention, Params>( + Params& params, cudaStream_t stream); +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_1.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..5fd75b0674ced9d94375a3bd5e4a9bfa8ab6a650 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_1.cu @@ -0,0 +1,11 @@ +#include +#include "single_prefill_sm90_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SingleFP8PrefillWithKVCacheDispatched + <128, MaskMode::kCausal, /*USE_SLIDING_WINDOW=*/false, DefaultFP8Attention, Params>( + Params& params, cudaStream_t stream); +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_2.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..b999545a84ba2d25e6a7e652cce4d09c5bf24d05 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_2.cu @@ -0,0 +1,11 @@ +#include +#include "single_prefill_sm90_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SingleFP8PrefillWithKVCacheDispatched + <128, MaskMode::kCustom, /*USE_SLIDING_WINDOW=*/false, DefaultFP8Attention, Params>( + Params& params, cudaStream_t stream); +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_3.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..88a253e381a7567fe45e244cb69307aa844a1953 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_3.cu @@ -0,0 +1,11 @@ +#include +#include "single_prefill_sm90_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SingleFP8PrefillWithKVCacheDispatched + <128, MaskMode::kMultiItemScoring, /*USE_SLIDING_WINDOW=*/false, DefaultFP8Attention, Params>( + Params& params, cudaStream_t stream); +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_fp8_sm90.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_fp8_sm90.cu new file mode 100644 index 0000000000000000000000000000000000000000..2e4eb9133649c0b339f2352e2e7ae49475c5c99a --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_fp8_sm90.cu @@ -0,0 +1,101 @@ +/* + * Copyright (c) 2024 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include +#include + +#include "pytorch_extension_utils.h" +#include "single_prefill_sm90_config.inc" + +namespace flashinfer { + +template +cudaError_t SingleFP8PrefillWithKVCacheDispatched(Params& params, cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +void single_prefill_with_kv_cache_sm90(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS) { + unsigned int head_dim_qk = q.size(2); + unsigned int head_dim_vo = v.size(2); + unsigned int num_qo_heads = q.size(1); + unsigned int qo_len = q.size(0); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + QKVLayout kv_layout = static_cast(layout); + const c10::cuda::OptionalCUDAGuard device_guard(q.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + const MaskMode mask_mode = static_cast(mask_mode_code); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, USE_SLIDING_WINDOW, + USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + Params params; + params.q_ptr = static_cast(q.data_ptr()); + params.k_ptr = static_cast(k.data_ptr()); + params.v_ptr = static_cast(v.data_ptr()); + params.o_ptr = static_cast(o.data_ptr()); + params.lse_ptr = maybe_lse ? (static_cast(maybe_lse->data_ptr())) : nullptr; + params.q_stride_n = q.stride(0); + params.q_stride_h = q.stride(1); + params.o_stride_n = o.stride(0); + params.o_stride_h = o.stride(1); + if (kv_layout == QKVLayout::kNHD) { + params.k_stride_n = k.stride(0); + params.k_stride_h = k.stride(1); + params.v_stride_n = v.stride(0); + params.v_stride_h = v.stride(1); + } else { + params.k_stride_h = k.stride(0); + params.k_stride_n = k.stride(1); + params.v_stride_h = v.stride(0); + params.v_stride_n = v.stride(1); + } + params.qo_len = q.size(0); + params.kv_len = k.size(0); + params.num_qo_heads = q.size(1); + params.num_kv_heads = k.size(1); + params.causal = mask_mode == MaskMode::kCausal; + params.group_size = params.num_qo_heads / params.num_kv_heads; + + // Note(Yilong): this should be checked on Python Side + // Only support window_left == 0 for now + params.window_left = window_left; + + // Note(Yilong): all quantization parameters are set in additional_params + ADDITIONAL_PARAMS_SETTER + + // Not support various head_dim for now + static_assert(HEAD_DIM_QK == HEAD_DIM_VO, "head_dim_qk and head_dim_vo should be the same"); + // Currently only support same quantization precision + static_assert(std::is_same_v); + + cudaError_t status = + SingleFP8PrefillWithKVCacheDispatched(params, stream); + TORCH_CHECK(status == cudaSuccess, "single_prefill_with_kv_cache_sm90 failed with error: " + + std::string(cudaGetErrorString(status))); + return true; + }); +} diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_config.inc b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..d6e598a6b4478fac0a8313778a0d888a77bba66a --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_config.inc @@ -0,0 +1,74 @@ +#pragma once +#include +#include +#include +#include +#include +#include +#include + + +#define ADDITIONAL_FUNC_PARAMS , at::Tensor scale_q, at::Tensor scale_k, at::Tensor scale_v, double sm_scale +#define ADDITIONAL_PARAMS_SETTER params.additional_params.scale_q = static_cast(scale_q.data_ptr()); \ +params.additional_params.scale_k = static_cast(scale_k.data_ptr()); \ +params.additional_params.scale_v = static_cast(scale_v.data_ptr()); \ +params.additional_params.sm_scale = sm_scale; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { using AttentionVariant = DefaultFP8Attention; __VA_ARGS__(); }) + +using namespace flashinfer; + +using DTypeQ = cutlass_dtype_t<__nv_fp8_e4m3>; +using DTypeKV = cutlass_dtype_t<__nv_fp8_e4m3>; +using DTypeO = cutlass_dtype_t; +using IdType = cutlass_dtype_t; + +constexpr int HEAD_DIM_QK = 256; +constexpr int HEAD_DIM_VO = 256; +constexpr auto USE_LOGITS_SOFT_CAP = true; +constexpr auto USE_SLIDING_WINDOW = true; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + // The QKV matrices. + DTypeQ* q_ptr; + DTypeKV* k_ptr; + DTypeKV* v_ptr; + DTypeO* o_ptr; + float* lse_ptr; + + // Additional params + struct AdditionalParams { + float* scale_q; +float* scale_k; +float* scale_v; +double sm_scale; +; + } additional_params; + + int64_t q_stride_n; + int64_t k_stride_n; + int64_t v_stride_n; + int64_t o_stride_n; + int64_t q_stride_h; + int64_t k_stride_h; + int64_t v_stride_h; + int64_t o_stride_h; + + int qo_len; + int kv_len; + int head_dim; + int num_qo_heads; + int num_kv_heads; + int group_size; + int window_left; + + bool causal; +}; + +#include \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_jit_pybind.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..2d6dd6076239658b22af333487f1f0fb49a86b1b --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_jit_pybind.cu @@ -0,0 +1,27 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "pytorch_extension_utils.h" +#include "single_prefill_sm90_config.inc" + +void single_prefill_with_kv_cache_sm90(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Single-request prefill attention with KV-Cache operator + m.def("run", single_prefill_with_kv_cache_sm90); +} diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_kernel_mask_0.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..458131cd472f1fced4c1915330817d6df7347289 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_kernel_mask_0.cu @@ -0,0 +1,11 @@ +#include +#include "single_prefill_sm90_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SingleFP8PrefillWithKVCacheDispatched + <256, MaskMode::kNone, /*USE_SLIDING_WINDOW=*/true, DefaultFP8Attention, Params>( + Params& params, cudaStream_t stream); +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_kernel_mask_1.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..b3c4557fa058a82676a8b950462c96dac1a2084c --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_kernel_mask_1.cu @@ -0,0 +1,11 @@ +#include +#include "single_prefill_sm90_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SingleFP8PrefillWithKVCacheDispatched + <256, MaskMode::kCausal, /*USE_SLIDING_WINDOW=*/true, DefaultFP8Attention, Params>( + Params& params, cudaStream_t stream); +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_kernel_mask_2.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..d53f13a489cfc6695ac2001f72d84a06a8a4f002 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_kernel_mask_2.cu @@ -0,0 +1,11 @@ +#include +#include "single_prefill_sm90_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SingleFP8PrefillWithKVCacheDispatched + <256, MaskMode::kCustom, /*USE_SLIDING_WINDOW=*/true, DefaultFP8Attention, Params>( + Params& params, cudaStream_t stream); +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_kernel_mask_3.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..87a1425f18477ef8ff6c5711253b40e5a9e1ed32 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_bf16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_kernel_mask_3.cu @@ -0,0 +1,11 @@ +#include +#include "single_prefill_sm90_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SingleFP8PrefillWithKVCacheDispatched + <256, MaskMode::kMultiItemScoring, /*USE_SLIDING_WINDOW=*/true, DefaultFP8Attention, Params>( + Params& params, cudaStream_t stream); +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_fp8_sm90.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_fp8_sm90.cu new file mode 100644 index 0000000000000000000000000000000000000000..2e4eb9133649c0b339f2352e2e7ae49475c5c99a --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_fp8_sm90.cu @@ -0,0 +1,101 @@ +/* + * Copyright (c) 2024 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include +#include + +#include "pytorch_extension_utils.h" +#include "single_prefill_sm90_config.inc" + +namespace flashinfer { + +template +cudaError_t SingleFP8PrefillWithKVCacheDispatched(Params& params, cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +void single_prefill_with_kv_cache_sm90(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS) { + unsigned int head_dim_qk = q.size(2); + unsigned int head_dim_vo = v.size(2); + unsigned int num_qo_heads = q.size(1); + unsigned int qo_len = q.size(0); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + QKVLayout kv_layout = static_cast(layout); + const c10::cuda::OptionalCUDAGuard device_guard(q.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + const MaskMode mask_mode = static_cast(mask_mode_code); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, USE_SLIDING_WINDOW, + USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + Params params; + params.q_ptr = static_cast(q.data_ptr()); + params.k_ptr = static_cast(k.data_ptr()); + params.v_ptr = static_cast(v.data_ptr()); + params.o_ptr = static_cast(o.data_ptr()); + params.lse_ptr = maybe_lse ? (static_cast(maybe_lse->data_ptr())) : nullptr; + params.q_stride_n = q.stride(0); + params.q_stride_h = q.stride(1); + params.o_stride_n = o.stride(0); + params.o_stride_h = o.stride(1); + if (kv_layout == QKVLayout::kNHD) { + params.k_stride_n = k.stride(0); + params.k_stride_h = k.stride(1); + params.v_stride_n = v.stride(0); + params.v_stride_h = v.stride(1); + } else { + params.k_stride_h = k.stride(0); + params.k_stride_n = k.stride(1); + params.v_stride_h = v.stride(0); + params.v_stride_n = v.stride(1); + } + params.qo_len = q.size(0); + params.kv_len = k.size(0); + params.num_qo_heads = q.size(1); + params.num_kv_heads = k.size(1); + params.causal = mask_mode == MaskMode::kCausal; + params.group_size = params.num_qo_heads / params.num_kv_heads; + + // Note(Yilong): this should be checked on Python Side + // Only support window_left == 0 for now + params.window_left = window_left; + + // Note(Yilong): all quantization parameters are set in additional_params + ADDITIONAL_PARAMS_SETTER + + // Not support various head_dim for now + static_assert(HEAD_DIM_QK == HEAD_DIM_VO, "head_dim_qk and head_dim_vo should be the same"); + // Currently only support same quantization precision + static_assert(std::is_same_v); + + cudaError_t status = + SingleFP8PrefillWithKVCacheDispatched(params, stream); + TORCH_CHECK(status == cudaSuccess, "single_prefill_with_kv_cache_sm90 failed with error: " + + std::string(cudaGetErrorString(status))); + return true; + }); +} diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_config.inc b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..8fd1d6c348080a76bbe5b7ad888fe12d36ecd504 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_config.inc @@ -0,0 +1,74 @@ +#pragma once +#include +#include +#include +#include +#include +#include +#include + + +#define ADDITIONAL_FUNC_PARAMS , at::Tensor scale_q, at::Tensor scale_k, at::Tensor scale_v, double sm_scale +#define ADDITIONAL_PARAMS_SETTER params.additional_params.scale_q = static_cast(scale_q.data_ptr()); \ +params.additional_params.scale_k = static_cast(scale_k.data_ptr()); \ +params.additional_params.scale_v = static_cast(scale_v.data_ptr()); \ +params.additional_params.sm_scale = sm_scale; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { using AttentionVariant = DefaultFP8Attention; __VA_ARGS__(); }) + +using namespace flashinfer; + +using DTypeQ = cutlass_dtype_t<__nv_fp8_e4m3>; +using DTypeKV = cutlass_dtype_t<__nv_fp8_e4m3>; +using DTypeO = cutlass_dtype_t; +using IdType = cutlass_dtype_t; + +constexpr int HEAD_DIM_QK = 128; +constexpr int HEAD_DIM_VO = 128; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto USE_SLIDING_WINDOW = false; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + // The QKV matrices. + DTypeQ* q_ptr; + DTypeKV* k_ptr; + DTypeKV* v_ptr; + DTypeO* o_ptr; + float* lse_ptr; + + // Additional params + struct AdditionalParams { + float* scale_q; +float* scale_k; +float* scale_v; +double sm_scale; +; + } additional_params; + + int64_t q_stride_n; + int64_t k_stride_n; + int64_t v_stride_n; + int64_t o_stride_n; + int64_t q_stride_h; + int64_t k_stride_h; + int64_t v_stride_h; + int64_t o_stride_h; + + int qo_len; + int kv_len; + int head_dim; + int num_qo_heads; + int num_kv_heads; + int group_size; + int window_left; + + bool causal; +}; + +#include \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_jit_pybind.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..2d6dd6076239658b22af333487f1f0fb49a86b1b --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_jit_pybind.cu @@ -0,0 +1,27 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "pytorch_extension_utils.h" +#include "single_prefill_sm90_config.inc" + +void single_prefill_with_kv_cache_sm90(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Single-request prefill attention with KV-Cache operator + m.def("run", single_prefill_with_kv_cache_sm90); +} diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_0.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..cea7f733ac8b56b095898da8c22d733aa53c7a3b --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_0.cu @@ -0,0 +1,11 @@ +#include +#include "single_prefill_sm90_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SingleFP8PrefillWithKVCacheDispatched + <128, MaskMode::kNone, /*USE_SLIDING_WINDOW=*/false, DefaultFP8Attention, Params>( + Params& params, cudaStream_t stream); +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_1.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..5fd75b0674ced9d94375a3bd5e4a9bfa8ab6a650 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_1.cu @@ -0,0 +1,11 @@ +#include +#include "single_prefill_sm90_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SingleFP8PrefillWithKVCacheDispatched + <128, MaskMode::kCausal, /*USE_SLIDING_WINDOW=*/false, DefaultFP8Attention, Params>( + Params& params, cudaStream_t stream); +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_2.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..b999545a84ba2d25e6a7e652cce4d09c5bf24d05 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_2.cu @@ -0,0 +1,11 @@ +#include +#include "single_prefill_sm90_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SingleFP8PrefillWithKVCacheDispatched + <128, MaskMode::kCustom, /*USE_SLIDING_WINDOW=*/false, DefaultFP8Attention, Params>( + Params& params, cudaStream_t stream); +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_3.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..88a253e381a7567fe45e244cb69307aa844a1953 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_3.cu @@ -0,0 +1,11 @@ +#include +#include "single_prefill_sm90_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SingleFP8PrefillWithKVCacheDispatched + <128, MaskMode::kMultiItemScoring, /*USE_SLIDING_WINDOW=*/false, DefaultFP8Attention, Params>( + Params& params, cudaStream_t stream); +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_fp8_sm90.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_fp8_sm90.cu new file mode 100644 index 0000000000000000000000000000000000000000..2e4eb9133649c0b339f2352e2e7ae49475c5c99a --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_fp8_sm90.cu @@ -0,0 +1,101 @@ +/* + * Copyright (c) 2024 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include +#include + +#include "pytorch_extension_utils.h" +#include "single_prefill_sm90_config.inc" + +namespace flashinfer { + +template +cudaError_t SingleFP8PrefillWithKVCacheDispatched(Params& params, cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +void single_prefill_with_kv_cache_sm90(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS) { + unsigned int head_dim_qk = q.size(2); + unsigned int head_dim_vo = v.size(2); + unsigned int num_qo_heads = q.size(1); + unsigned int qo_len = q.size(0); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + QKVLayout kv_layout = static_cast(layout); + const c10::cuda::OptionalCUDAGuard device_guard(q.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + const MaskMode mask_mode = static_cast(mask_mode_code); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, USE_SLIDING_WINDOW, + USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + Params params; + params.q_ptr = static_cast(q.data_ptr()); + params.k_ptr = static_cast(k.data_ptr()); + params.v_ptr = static_cast(v.data_ptr()); + params.o_ptr = static_cast(o.data_ptr()); + params.lse_ptr = maybe_lse ? (static_cast(maybe_lse->data_ptr())) : nullptr; + params.q_stride_n = q.stride(0); + params.q_stride_h = q.stride(1); + params.o_stride_n = o.stride(0); + params.o_stride_h = o.stride(1); + if (kv_layout == QKVLayout::kNHD) { + params.k_stride_n = k.stride(0); + params.k_stride_h = k.stride(1); + params.v_stride_n = v.stride(0); + params.v_stride_h = v.stride(1); + } else { + params.k_stride_h = k.stride(0); + params.k_stride_n = k.stride(1); + params.v_stride_h = v.stride(0); + params.v_stride_n = v.stride(1); + } + params.qo_len = q.size(0); + params.kv_len = k.size(0); + params.num_qo_heads = q.size(1); + params.num_kv_heads = k.size(1); + params.causal = mask_mode == MaskMode::kCausal; + params.group_size = params.num_qo_heads / params.num_kv_heads; + + // Note(Yilong): this should be checked on Python Side + // Only support window_left == 0 for now + params.window_left = window_left; + + // Note(Yilong): all quantization parameters are set in additional_params + ADDITIONAL_PARAMS_SETTER + + // Not support various head_dim for now + static_assert(HEAD_DIM_QK == HEAD_DIM_VO, "head_dim_qk and head_dim_vo should be the same"); + // Currently only support same quantization precision + static_assert(std::is_same_v); + + cudaError_t status = + SingleFP8PrefillWithKVCacheDispatched(params, stream); + TORCH_CHECK(status == cudaSuccess, "single_prefill_with_kv_cache_sm90 failed with error: " + + std::string(cudaGetErrorString(status))); + return true; + }); +} diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_config.inc b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..7ab1ecb8ff022a69cf2122bf32fb3aa07ded9cd7 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_config.inc @@ -0,0 +1,74 @@ +#pragma once +#include +#include +#include +#include +#include +#include +#include + + +#define ADDITIONAL_FUNC_PARAMS , at::Tensor scale_q, at::Tensor scale_k, at::Tensor scale_v, double sm_scale +#define ADDITIONAL_PARAMS_SETTER params.additional_params.scale_q = static_cast(scale_q.data_ptr()); \ +params.additional_params.scale_k = static_cast(scale_k.data_ptr()); \ +params.additional_params.scale_v = static_cast(scale_v.data_ptr()); \ +params.additional_params.sm_scale = sm_scale; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { using AttentionVariant = DefaultFP8Attention; __VA_ARGS__(); }) + +using namespace flashinfer; + +using DTypeQ = cutlass_dtype_t<__nv_fp8_e4m3>; +using DTypeKV = cutlass_dtype_t<__nv_fp8_e4m3>; +using DTypeO = cutlass_dtype_t; +using IdType = cutlass_dtype_t; + +constexpr int HEAD_DIM_QK = 256; +constexpr int HEAD_DIM_VO = 256; +constexpr auto USE_LOGITS_SOFT_CAP = true; +constexpr auto USE_SLIDING_WINDOW = true; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + // The QKV matrices. + DTypeQ* q_ptr; + DTypeKV* k_ptr; + DTypeKV* v_ptr; + DTypeO* o_ptr; + float* lse_ptr; + + // Additional params + struct AdditionalParams { + float* scale_q; +float* scale_k; +float* scale_v; +double sm_scale; +; + } additional_params; + + int64_t q_stride_n; + int64_t k_stride_n; + int64_t v_stride_n; + int64_t o_stride_n; + int64_t q_stride_h; + int64_t k_stride_h; + int64_t v_stride_h; + int64_t o_stride_h; + + int qo_len; + int kv_len; + int head_dim; + int num_qo_heads; + int num_kv_heads; + int group_size; + int window_left; + + bool causal; +}; + +#include \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_jit_pybind.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..2d6dd6076239658b22af333487f1f0fb49a86b1b --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_jit_pybind.cu @@ -0,0 +1,27 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "pytorch_extension_utils.h" +#include "single_prefill_sm90_config.inc" + +void single_prefill_with_kv_cache_sm90(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Single-request prefill attention with KV-Cache operator + m.def("run", single_prefill_with_kv_cache_sm90); +} diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_kernel_mask_0.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..458131cd472f1fced4c1915330817d6df7347289 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_kernel_mask_0.cu @@ -0,0 +1,11 @@ +#include +#include "single_prefill_sm90_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SingleFP8PrefillWithKVCacheDispatched + <256, MaskMode::kNone, /*USE_SLIDING_WINDOW=*/true, DefaultFP8Attention, Params>( + Params& params, cudaStream_t stream); +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_kernel_mask_1.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..b3c4557fa058a82676a8b950462c96dac1a2084c --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_kernel_mask_1.cu @@ -0,0 +1,11 @@ +#include +#include "single_prefill_sm90_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SingleFP8PrefillWithKVCacheDispatched + <256, MaskMode::kCausal, /*USE_SLIDING_WINDOW=*/true, DefaultFP8Attention, Params>( + Params& params, cudaStream_t stream); +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_kernel_mask_2.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..d53f13a489cfc6695ac2001f72d84a06a8a4f002 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_kernel_mask_2.cu @@ -0,0 +1,11 @@ +#include +#include "single_prefill_sm90_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SingleFP8PrefillWithKVCacheDispatched + <256, MaskMode::kCustom, /*USE_SLIDING_WINDOW=*/true, DefaultFP8Attention, Params>( + Params& params, cudaStream_t stream); +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_kernel_mask_3.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..87a1425f18477ef8ff6c5711253b40e5a9e1ed32 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_e4m3_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_kernel_mask_3.cu @@ -0,0 +1,11 @@ +#include +#include "single_prefill_sm90_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SingleFP8PrefillWithKVCacheDispatched + <256, MaskMode::kMultiItemScoring, /*USE_SLIDING_WINDOW=*/true, DefaultFP8Attention, Params>( + Params& params, cudaStream_t stream); +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill.cu new file mode 100644 index 0000000000000000000000000000000000000000..8d77f8420680e70a7929d8b16bd95f5587a01752 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill.cu @@ -0,0 +1,114 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include + +#include "flashinfer/fastdiv.cuh" +#include "pytorch_extension_utils.h" +#include "single_prefill_config.inc" + +namespace flashinfer { + +template +cudaError_t SinglePrefillWithKVCacheDispatched(Params params, typename Params::DTypeO* tmp, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +void single_prefill_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS) { + auto device = q.device(); + unsigned int head_dim_qk = q.size(2); + unsigned int kv_len, qo_len, num_kv_heads, num_qo_heads; + QKVLayout kv_layout = static_cast(layout); + qo_len = q.size(0); + num_qo_heads = q.size(1); + uint32_t q_stride_n = q.stride(0), q_stride_h = q.stride(1), k_stride_n, k_stride_h, v_stride_n, + v_stride_h; + if (kv_layout == QKVLayout::kNHD) { + kv_len = k.size(0); + num_kv_heads = k.size(1); + k_stride_n = k.stride(0); + k_stride_h = k.stride(1); + v_stride_n = v.stride(0); + v_stride_h = v.stride(1); + } else { + kv_len = k.size(1); + num_kv_heads = k.size(0); + k_stride_h = k.stride(0); + k_stride_n = k.stride(1); + v_stride_h = v.stride(0); + v_stride_n = v.stride(1); + } + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == qo_len, lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == num_qo_heads, lse.size(1), q.size(1)); + } + + const MaskMode mask_mode = static_cast(mask_mode_code); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(device); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, Params, + [&] { + Params params; + + params.q = static_cast(q.data_ptr()); + params.k = static_cast(k.data_ptr()); + params.v = static_cast(v.data_ptr()); + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.num_qo_heads = num_qo_heads; + params.num_kv_heads = num_kv_heads; + params.group_size = uint_fastdiv(num_qo_heads / num_kv_heads); + params.qo_len = qo_len; + params.kv_len = kv_len; + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.k_stride_n = k_stride_n; + params.k_stride_h = k_stride_h; + params.v_stride_n = v_stride_n; + params.v_stride_h = v_stride_h; + + params.window_left = window_left; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + cudaError_t status = flashinfer::SinglePrefillWithKVCacheDispatched< + HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + /*use_fp16_qk_reduction=*/USE_FP16_QK_REDUCTION, MASK_MODE, AttentionVariant>( + params, static_cast(tmp.data_ptr()), stream); + TORCH_CHECK(status == cudaSuccess, + "SinglePrefillWithKVCache kernel launch failed, error: " + + std::string(cudaGetErrorString(status))); + return true; + }); +} diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_config.inc b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..1bb436b42992f6f5e9a74fbe9cea362d9cdca170 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_config.inc @@ -0,0 +1,83 @@ +#pragma once +#include +#include +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_custom_mask, std::optional maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta +#define ADDITIONAL_PARAMS_SETTER params.maybe_custom_mask = maybe_custom_mask ? static_cast(maybe_custom_mask->data_ptr()): nullptr; \ +params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; + + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, Params, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { \ + constexpr bool use_custom_mask = MASK_MODE == MaskMode::kCustom; \ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ + }) + + +using namespace flashinfer; + +using DTypeQ = half; +using DTypeKV = __nv_fp8_e4m3; +using DTypeO = half; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 128; +constexpr int HEAD_DIM_VO = 128; +constexpr bool USE_FP16_QK_REDUCTION = false; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = false; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = int32_t; + DTypeQ* q; + DTypeKV* k; + DTypeKV* v; + DTypeO* o; + float* lse; + uint_fastdiv group_size; + + uint8_t* maybe_custom_mask; +float* maybe_alibi_slopes; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; + + + uint32_t qo_len; + uint32_t kv_len; + uint32_t num_qo_heads; + uint32_t num_kv_heads; + uint32_t q_stride_n; + uint32_t q_stride_h; + uint32_t k_stride_n; + uint32_t k_stride_h; + uint32_t v_stride_n; + uint32_t v_stride_h; + uint32_t head_dim; + int32_t window_left; + + bool partition_kv; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { + return qo_len; + } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return kv_len; + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_jit_pybind.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..aaccbc045f9fd8566654a106379c99a053489fc7 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_jit_pybind.cu @@ -0,0 +1,27 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "pytorch_extension_utils.h" +#include "single_prefill_config.inc" + +void single_prefill_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Single-request prefill attention with KV-Cache operator + m.def("run", single_prefill_with_kv_cache); +} diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_0.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..71281e7ec1ed88578741aa75fe975b0569200d92 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_0.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kNone == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 128, 128, PosEncodingMode::kNone, false, MaskMode::kNone, DefaultAttention, Params>( + Params params, half* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_1.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..856c9b6a04da54a98ee6a1cfea766c53a49dcae8 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_1.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCausal == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 128, 128, PosEncodingMode::kNone, false, MaskMode::kCausal, DefaultAttention, Params>( + Params params, half* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_2.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..11bc10741f0ce80f34d38bc9079695a140a43d3d --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_2.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCustom == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 128, 128, PosEncodingMode::kNone, false, MaskMode::kCustom, DefaultAttention, Params>( + Params params, half* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_3.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..8b6097619e0aab70a83d486e74db0d12ba6e7b7b --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_3.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kMultiItemScoring == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 128, 128, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, DefaultAttention, Params>( + Params params, half* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill.cu new file mode 100644 index 0000000000000000000000000000000000000000..8d77f8420680e70a7929d8b16bd95f5587a01752 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill.cu @@ -0,0 +1,114 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include + +#include "flashinfer/fastdiv.cuh" +#include "pytorch_extension_utils.h" +#include "single_prefill_config.inc" + +namespace flashinfer { + +template +cudaError_t SinglePrefillWithKVCacheDispatched(Params params, typename Params::DTypeO* tmp, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +void single_prefill_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS) { + auto device = q.device(); + unsigned int head_dim_qk = q.size(2); + unsigned int kv_len, qo_len, num_kv_heads, num_qo_heads; + QKVLayout kv_layout = static_cast(layout); + qo_len = q.size(0); + num_qo_heads = q.size(1); + uint32_t q_stride_n = q.stride(0), q_stride_h = q.stride(1), k_stride_n, k_stride_h, v_stride_n, + v_stride_h; + if (kv_layout == QKVLayout::kNHD) { + kv_len = k.size(0); + num_kv_heads = k.size(1); + k_stride_n = k.stride(0); + k_stride_h = k.stride(1); + v_stride_n = v.stride(0); + v_stride_h = v.stride(1); + } else { + kv_len = k.size(1); + num_kv_heads = k.size(0); + k_stride_h = k.stride(0); + k_stride_n = k.stride(1); + v_stride_h = v.stride(0); + v_stride_n = v.stride(1); + } + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == qo_len, lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == num_qo_heads, lse.size(1), q.size(1)); + } + + const MaskMode mask_mode = static_cast(mask_mode_code); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(device); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, Params, + [&] { + Params params; + + params.q = static_cast(q.data_ptr()); + params.k = static_cast(k.data_ptr()); + params.v = static_cast(v.data_ptr()); + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.num_qo_heads = num_qo_heads; + params.num_kv_heads = num_kv_heads; + params.group_size = uint_fastdiv(num_qo_heads / num_kv_heads); + params.qo_len = qo_len; + params.kv_len = kv_len; + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.k_stride_n = k_stride_n; + params.k_stride_h = k_stride_h; + params.v_stride_n = v_stride_n; + params.v_stride_h = v_stride_h; + + params.window_left = window_left; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + cudaError_t status = flashinfer::SinglePrefillWithKVCacheDispatched< + HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + /*use_fp16_qk_reduction=*/USE_FP16_QK_REDUCTION, MASK_MODE, AttentionVariant>( + params, static_cast(tmp.data_ptr()), stream); + TORCH_CHECK(status == cudaSuccess, + "SinglePrefillWithKVCache kernel launch failed, error: " + + std::string(cudaGetErrorString(status))); + return true; + }); +} diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_config.inc b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..a3a9e9d16589bb41d91425af8173c6f78b3490ca --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_config.inc @@ -0,0 +1,83 @@ +#pragma once +#include +#include +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_custom_mask, std::optional maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta +#define ADDITIONAL_PARAMS_SETTER params.maybe_custom_mask = maybe_custom_mask ? static_cast(maybe_custom_mask->data_ptr()): nullptr; \ +params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; + + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, Params, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { \ + constexpr bool use_custom_mask = MASK_MODE == MaskMode::kCustom; \ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ + }) + + +using namespace flashinfer; + +using DTypeQ = half; +using DTypeKV = __nv_fp8_e4m3; +using DTypeO = half; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 256; +constexpr int HEAD_DIM_VO = 256; +constexpr bool USE_FP16_QK_REDUCTION = false; +constexpr auto USE_LOGITS_SOFT_CAP = true; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = true; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = int32_t; + DTypeQ* q; + DTypeKV* k; + DTypeKV* v; + DTypeO* o; + float* lse; + uint_fastdiv group_size; + + uint8_t* maybe_custom_mask; +float* maybe_alibi_slopes; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; + + + uint32_t qo_len; + uint32_t kv_len; + uint32_t num_qo_heads; + uint32_t num_kv_heads; + uint32_t q_stride_n; + uint32_t q_stride_h; + uint32_t k_stride_n; + uint32_t k_stride_h; + uint32_t v_stride_n; + uint32_t v_stride_h; + uint32_t head_dim; + int32_t window_left; + + bool partition_kv; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { + return qo_len; + } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return kv_len; + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_jit_pybind.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..aaccbc045f9fd8566654a106379c99a053489fc7 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_jit_pybind.cu @@ -0,0 +1,27 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "pytorch_extension_utils.h" +#include "single_prefill_config.inc" + +void single_prefill_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Single-request prefill attention with KV-Cache operator + m.def("run", single_prefill_with_kv_cache); +} diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_kernel_mask_0.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..14843c3fb10f36c123129d10410958397c0add42 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_kernel_mask_0.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kNone == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 256, 256, PosEncodingMode::kNone, false, MaskMode::kNone, DefaultAttention, Params>( + Params params, half* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_kernel_mask_1.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..dd1c37887e9c496433beabd47a6c0e627a6a5468 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_kernel_mask_1.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCausal == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 256, 256, PosEncodingMode::kNone, false, MaskMode::kCausal, DefaultAttention, Params>( + Params params, half* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_kernel_mask_2.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..486980ebafeeba7400fc788925911ccb7337400e --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_kernel_mask_2.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCustom == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 256, 256, PosEncodingMode::kNone, false, MaskMode::kCustom, DefaultAttention, Params>( + Params params, half* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_kernel_mask_3.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..a4c568fac4ff8089c8e50ec27b322291cabf84da --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_kernel_mask_3.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kMultiItemScoring == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 256, 256, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, DefaultAttention, Params>( + Params params, half* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill.cu new file mode 100644 index 0000000000000000000000000000000000000000..8d77f8420680e70a7929d8b16bd95f5587a01752 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill.cu @@ -0,0 +1,114 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include + +#include "flashinfer/fastdiv.cuh" +#include "pytorch_extension_utils.h" +#include "single_prefill_config.inc" + +namespace flashinfer { + +template +cudaError_t SinglePrefillWithKVCacheDispatched(Params params, typename Params::DTypeO* tmp, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +void single_prefill_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS) { + auto device = q.device(); + unsigned int head_dim_qk = q.size(2); + unsigned int kv_len, qo_len, num_kv_heads, num_qo_heads; + QKVLayout kv_layout = static_cast(layout); + qo_len = q.size(0); + num_qo_heads = q.size(1); + uint32_t q_stride_n = q.stride(0), q_stride_h = q.stride(1), k_stride_n, k_stride_h, v_stride_n, + v_stride_h; + if (kv_layout == QKVLayout::kNHD) { + kv_len = k.size(0); + num_kv_heads = k.size(1); + k_stride_n = k.stride(0); + k_stride_h = k.stride(1); + v_stride_n = v.stride(0); + v_stride_h = v.stride(1); + } else { + kv_len = k.size(1); + num_kv_heads = k.size(0); + k_stride_h = k.stride(0); + k_stride_n = k.stride(1); + v_stride_h = v.stride(0); + v_stride_n = v.stride(1); + } + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == qo_len, lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == num_qo_heads, lse.size(1), q.size(1)); + } + + const MaskMode mask_mode = static_cast(mask_mode_code); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(device); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, Params, + [&] { + Params params; + + params.q = static_cast(q.data_ptr()); + params.k = static_cast(k.data_ptr()); + params.v = static_cast(v.data_ptr()); + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.num_qo_heads = num_qo_heads; + params.num_kv_heads = num_kv_heads; + params.group_size = uint_fastdiv(num_qo_heads / num_kv_heads); + params.qo_len = qo_len; + params.kv_len = kv_len; + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.k_stride_n = k_stride_n; + params.k_stride_h = k_stride_h; + params.v_stride_n = v_stride_n; + params.v_stride_h = v_stride_h; + + params.window_left = window_left; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + cudaError_t status = flashinfer::SinglePrefillWithKVCacheDispatched< + HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + /*use_fp16_qk_reduction=*/USE_FP16_QK_REDUCTION, MASK_MODE, AttentionVariant>( + params, static_cast(tmp.data_ptr()), stream); + TORCH_CHECK(status == cudaSuccess, + "SinglePrefillWithKVCache kernel launch failed, error: " + + std::string(cudaGetErrorString(status))); + return true; + }); +} diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_config.inc b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..7f541ab69be2381f5665addcc7b52e374fa3d1e2 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_config.inc @@ -0,0 +1,83 @@ +#pragma once +#include +#include +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_custom_mask, std::optional maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta +#define ADDITIONAL_PARAMS_SETTER params.maybe_custom_mask = maybe_custom_mask ? static_cast(maybe_custom_mask->data_ptr()): nullptr; \ +params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; + + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, Params, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { \ + constexpr bool use_custom_mask = MASK_MODE == MaskMode::kCustom; \ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ + }) + + +using namespace flashinfer; + +using DTypeQ = half; +using DTypeKV = __nv_fp8_e4m3; +using DTypeO = half; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 64; +constexpr int HEAD_DIM_VO = 64; +constexpr bool USE_FP16_QK_REDUCTION = false; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = false; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = int32_t; + DTypeQ* q; + DTypeKV* k; + DTypeKV* v; + DTypeO* o; + float* lse; + uint_fastdiv group_size; + + uint8_t* maybe_custom_mask; +float* maybe_alibi_slopes; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; + + + uint32_t qo_len; + uint32_t kv_len; + uint32_t num_qo_heads; + uint32_t num_kv_heads; + uint32_t q_stride_n; + uint32_t q_stride_h; + uint32_t k_stride_n; + uint32_t k_stride_h; + uint32_t v_stride_n; + uint32_t v_stride_h; + uint32_t head_dim; + int32_t window_left; + + bool partition_kv; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { + return qo_len; + } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return kv_len; + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_jit_pybind.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..aaccbc045f9fd8566654a106379c99a053489fc7 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_jit_pybind.cu @@ -0,0 +1,27 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "pytorch_extension_utils.h" +#include "single_prefill_config.inc" + +void single_prefill_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Single-request prefill attention with KV-Cache operator + m.def("run", single_prefill_with_kv_cache); +} diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_0.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..465ba8aa40f1852772e15b84eace44711f34a177 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_0.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kNone == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, DefaultAttention, Params>( + Params params, half* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_1.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..1fd3d0a416bf7cfacd861519ac03ed387b5b6526 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_1.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCausal == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, DefaultAttention, Params>( + Params params, half* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_2.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..f71cb040b06f03780749c54b00cbae78a3872b68 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_2.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCustom == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, DefaultAttention, Params>( + Params params, half* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_3.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..6800c5a9275289669e7107ded5bb7bc13245b5f6 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_3.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kMultiItemScoring == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, DefaultAttention, Params>( + Params params, half* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill.cu new file mode 100644 index 0000000000000000000000000000000000000000..8d77f8420680e70a7929d8b16bd95f5587a01752 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill.cu @@ -0,0 +1,114 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include + +#include "flashinfer/fastdiv.cuh" +#include "pytorch_extension_utils.h" +#include "single_prefill_config.inc" + +namespace flashinfer { + +template +cudaError_t SinglePrefillWithKVCacheDispatched(Params params, typename Params::DTypeO* tmp, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +void single_prefill_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS) { + auto device = q.device(); + unsigned int head_dim_qk = q.size(2); + unsigned int kv_len, qo_len, num_kv_heads, num_qo_heads; + QKVLayout kv_layout = static_cast(layout); + qo_len = q.size(0); + num_qo_heads = q.size(1); + uint32_t q_stride_n = q.stride(0), q_stride_h = q.stride(1), k_stride_n, k_stride_h, v_stride_n, + v_stride_h; + if (kv_layout == QKVLayout::kNHD) { + kv_len = k.size(0); + num_kv_heads = k.size(1); + k_stride_n = k.stride(0); + k_stride_h = k.stride(1); + v_stride_n = v.stride(0); + v_stride_h = v.stride(1); + } else { + kv_len = k.size(1); + num_kv_heads = k.size(0); + k_stride_h = k.stride(0); + k_stride_n = k.stride(1); + v_stride_h = v.stride(0); + v_stride_n = v.stride(1); + } + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == qo_len, lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == num_qo_heads, lse.size(1), q.size(1)); + } + + const MaskMode mask_mode = static_cast(mask_mode_code); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(device); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, Params, + [&] { + Params params; + + params.q = static_cast(q.data_ptr()); + params.k = static_cast(k.data_ptr()); + params.v = static_cast(v.data_ptr()); + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.num_qo_heads = num_qo_heads; + params.num_kv_heads = num_kv_heads; + params.group_size = uint_fastdiv(num_qo_heads / num_kv_heads); + params.qo_len = qo_len; + params.kv_len = kv_len; + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.k_stride_n = k_stride_n; + params.k_stride_h = k_stride_h; + params.v_stride_n = v_stride_n; + params.v_stride_h = v_stride_h; + + params.window_left = window_left; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + cudaError_t status = flashinfer::SinglePrefillWithKVCacheDispatched< + HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + /*use_fp16_qk_reduction=*/USE_FP16_QK_REDUCTION, MASK_MODE, AttentionVariant>( + params, static_cast(tmp.data_ptr()), stream); + TORCH_CHECK(status == cudaSuccess, + "SinglePrefillWithKVCache kernel launch failed, error: " + + std::string(cudaGetErrorString(status))); + return true; + }); +} diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_config.inc b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..7e0999889d1ba6257083bbf765d560d4c46ce73e --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_config.inc @@ -0,0 +1,83 @@ +#pragma once +#include +#include +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_custom_mask, std::optional maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta +#define ADDITIONAL_PARAMS_SETTER params.maybe_custom_mask = maybe_custom_mask ? static_cast(maybe_custom_mask->data_ptr()): nullptr; \ +params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; + + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, Params, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { \ + constexpr bool use_custom_mask = MASK_MODE == MaskMode::kCustom; \ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ + }) + + +using namespace flashinfer; + +using DTypeQ = half; +using DTypeKV = __nv_fp8_e4m3; +using DTypeO = half; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 64; +constexpr int HEAD_DIM_VO = 64; +constexpr bool USE_FP16_QK_REDUCTION = false; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = true; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = int32_t; + DTypeQ* q; + DTypeKV* k; + DTypeKV* v; + DTypeO* o; + float* lse; + uint_fastdiv group_size; + + uint8_t* maybe_custom_mask; +float* maybe_alibi_slopes; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; + + + uint32_t qo_len; + uint32_t kv_len; + uint32_t num_qo_heads; + uint32_t num_kv_heads; + uint32_t q_stride_n; + uint32_t q_stride_h; + uint32_t k_stride_n; + uint32_t k_stride_h; + uint32_t v_stride_n; + uint32_t v_stride_h; + uint32_t head_dim; + int32_t window_left; + + bool partition_kv; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { + return qo_len; + } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return kv_len; + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_jit_pybind.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..aaccbc045f9fd8566654a106379c99a053489fc7 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_jit_pybind.cu @@ -0,0 +1,27 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "pytorch_extension_utils.h" +#include "single_prefill_config.inc" + +void single_prefill_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Single-request prefill attention with KV-Cache operator + m.def("run", single_prefill_with_kv_cache); +} diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_0.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..105f200cc3715c7f935961ea39f055f9f7a12b40 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_0.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kNone == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, DefaultAttention, Params>( + Params params, half* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_1.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..c835246808e83388370fd3b22b6703a5db464580 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_1.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCausal == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, DefaultAttention, Params>( + Params params, half* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_2.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..8104e776733a3ed680359f01197b9a172ac43248 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_2.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCustom == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, DefaultAttention, Params>( + Params params, half* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_3.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..b30f6e965634630a057766e440ffeb2ebe9cebd9 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_e4m3_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_3.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kMultiItemScoring == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, DefaultAttention, Params>( + Params params, half* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill.cu new file mode 100644 index 0000000000000000000000000000000000000000..8d77f8420680e70a7929d8b16bd95f5587a01752 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill.cu @@ -0,0 +1,114 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include + +#include "flashinfer/fastdiv.cuh" +#include "pytorch_extension_utils.h" +#include "single_prefill_config.inc" + +namespace flashinfer { + +template +cudaError_t SinglePrefillWithKVCacheDispatched(Params params, typename Params::DTypeO* tmp, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +void single_prefill_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS) { + auto device = q.device(); + unsigned int head_dim_qk = q.size(2); + unsigned int kv_len, qo_len, num_kv_heads, num_qo_heads; + QKVLayout kv_layout = static_cast(layout); + qo_len = q.size(0); + num_qo_heads = q.size(1); + uint32_t q_stride_n = q.stride(0), q_stride_h = q.stride(1), k_stride_n, k_stride_h, v_stride_n, + v_stride_h; + if (kv_layout == QKVLayout::kNHD) { + kv_len = k.size(0); + num_kv_heads = k.size(1); + k_stride_n = k.stride(0); + k_stride_h = k.stride(1); + v_stride_n = v.stride(0); + v_stride_h = v.stride(1); + } else { + kv_len = k.size(1); + num_kv_heads = k.size(0); + k_stride_h = k.stride(0); + k_stride_n = k.stride(1); + v_stride_h = v.stride(0); + v_stride_n = v.stride(1); + } + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == qo_len, lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == num_qo_heads, lse.size(1), q.size(1)); + } + + const MaskMode mask_mode = static_cast(mask_mode_code); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(device); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, Params, + [&] { + Params params; + + params.q = static_cast(q.data_ptr()); + params.k = static_cast(k.data_ptr()); + params.v = static_cast(v.data_ptr()); + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.num_qo_heads = num_qo_heads; + params.num_kv_heads = num_kv_heads; + params.group_size = uint_fastdiv(num_qo_heads / num_kv_heads); + params.qo_len = qo_len; + params.kv_len = kv_len; + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.k_stride_n = k_stride_n; + params.k_stride_h = k_stride_h; + params.v_stride_n = v_stride_n; + params.v_stride_h = v_stride_h; + + params.window_left = window_left; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + cudaError_t status = flashinfer::SinglePrefillWithKVCacheDispatched< + HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + /*use_fp16_qk_reduction=*/USE_FP16_QK_REDUCTION, MASK_MODE, AttentionVariant>( + params, static_cast(tmp.data_ptr()), stream); + TORCH_CHECK(status == cudaSuccess, + "SinglePrefillWithKVCache kernel launch failed, error: " + + std::string(cudaGetErrorString(status))); + return true; + }); +} diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_config.inc b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..fe16189ff8224a5edd61b6eb51061c109554d0aa --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_config.inc @@ -0,0 +1,83 @@ +#pragma once +#include +#include +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_custom_mask, std::optional maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta +#define ADDITIONAL_PARAMS_SETTER params.maybe_custom_mask = maybe_custom_mask ? static_cast(maybe_custom_mask->data_ptr()): nullptr; \ +params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; + + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, Params, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { \ + constexpr bool use_custom_mask = MASK_MODE == MaskMode::kCustom; \ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ + }) + + +using namespace flashinfer; + +using DTypeQ = half; +using DTypeKV = half; +using DTypeO = half; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 128; +constexpr int HEAD_DIM_VO = 128; +constexpr bool USE_FP16_QK_REDUCTION = false; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = false; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = int32_t; + DTypeQ* q; + DTypeKV* k; + DTypeKV* v; + DTypeO* o; + float* lse; + uint_fastdiv group_size; + + uint8_t* maybe_custom_mask; +float* maybe_alibi_slopes; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; + + + uint32_t qo_len; + uint32_t kv_len; + uint32_t num_qo_heads; + uint32_t num_kv_heads; + uint32_t q_stride_n; + uint32_t q_stride_h; + uint32_t k_stride_n; + uint32_t k_stride_h; + uint32_t v_stride_n; + uint32_t v_stride_h; + uint32_t head_dim; + int32_t window_left; + + bool partition_kv; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { + return qo_len; + } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return kv_len; + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_jit_pybind.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..aaccbc045f9fd8566654a106379c99a053489fc7 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_jit_pybind.cu @@ -0,0 +1,27 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "pytorch_extension_utils.h" +#include "single_prefill_config.inc" + +void single_prefill_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Single-request prefill attention with KV-Cache operator + m.def("run", single_prefill_with_kv_cache); +} diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_0.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..71281e7ec1ed88578741aa75fe975b0569200d92 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_0.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kNone == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 128, 128, PosEncodingMode::kNone, false, MaskMode::kNone, DefaultAttention, Params>( + Params params, half* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_1.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..856c9b6a04da54a98ee6a1cfea766c53a49dcae8 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_1.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCausal == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 128, 128, PosEncodingMode::kNone, false, MaskMode::kCausal, DefaultAttention, Params>( + Params params, half* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_2.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..11bc10741f0ce80f34d38bc9079695a140a43d3d --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_2.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCustom == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 128, 128, PosEncodingMode::kNone, false, MaskMode::kCustom, DefaultAttention, Params>( + Params params, half* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_3.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..8b6097619e0aab70a83d486e74db0d12ba6e7b7b --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_3.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kMultiItemScoring == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 128, 128, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, DefaultAttention, Params>( + Params params, half* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90.cu new file mode 100644 index 0000000000000000000000000000000000000000..39cabb9c2873d8e3922f6618e7081df62cac1b57 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90.cu @@ -0,0 +1,93 @@ +/* + * Copyright (c) 2024 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include +#include + +#include "pytorch_extension_utils.h" +#include "single_prefill_sm90_config.inc" + +namespace flashinfer { + +template +cudaError_t SinglePrefillWithKVCacheDispatched(Params& params, cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +void single_prefill_with_kv_cache_sm90(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS) { + unsigned int head_dim_qk = q.size(2); + unsigned int head_dim_vo = v.size(2); + unsigned int num_qo_heads = q.size(1); + unsigned int qo_len = q.size(0); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + QKVLayout kv_layout = static_cast(layout); + const c10::cuda::OptionalCUDAGuard device_guard(q.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + const MaskMode mask_mode = static_cast(mask_mode_code); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, USE_SLIDING_WINDOW, + USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + Params params; + params.q_ptr = static_cast(q.data_ptr()); + params.k_ptr = static_cast(k.data_ptr()); + params.v_ptr = static_cast(v.data_ptr()); + params.o_ptr = static_cast(o.data_ptr()); + params.lse_ptr = maybe_lse ? (static_cast(maybe_lse->data_ptr())) : nullptr; + params.q_stride_n = q.stride(0); + params.q_stride_h = q.stride(1); + params.o_stride_n = o.stride(0); + params.o_stride_h = o.stride(1); + if (kv_layout == QKVLayout::kNHD) { + params.k_stride_n = k.stride(0); + params.k_stride_h = k.stride(1); + params.v_stride_n = v.stride(0); + params.v_stride_h = v.stride(1); + } else { + params.k_stride_h = k.stride(0); + params.k_stride_n = k.stride(1); + params.v_stride_h = v.stride(0); + params.v_stride_n = v.stride(1); + } + params.qo_len = q.size(0); + params.kv_len = k.size(0); + params.num_qo_heads = q.size(1); + params.num_kv_heads = k.size(1); + params.causal = mask_mode == MaskMode::kCausal; + params.group_size = params.num_qo_heads / params.num_kv_heads; + params.window_left = window_left; + + ADDITIONAL_PARAMS_SETTER + + cudaError_t status = + SinglePrefillWithKVCacheDispatched(params, + stream); + TORCH_CHECK(status == cudaSuccess, "single_prefill_with_kv_cache_sm90 failed with error: " + + std::string(cudaGetErrorString(status))); + return true; + }); +} diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_config.inc b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..6c01f360b88fa5c9ef46a5625267df753da562fb --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_config.inc @@ -0,0 +1,70 @@ +#pragma once +#include +#include +#include +#include +#include +#include +#include + + +#define ADDITIONAL_FUNC_PARAMS , double logits_soft_cap, double sm_scale +#define ADDITIONAL_PARAMS_SETTER params.additional_params.logits_soft_cap = logits_soft_cap; \ +params.additional_params.sm_scale = sm_scale; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { using AttentionVariant = DefaultAttention; __VA_ARGS__(); }) + +using namespace flashinfer; + +using DTypeQ = cutlass_dtype_t; +using DTypeKV = cutlass_dtype_t; +using DTypeO = cutlass_dtype_t; +using IdType = cutlass_dtype_t; + +constexpr int HEAD_DIM_QK = 128; +constexpr int HEAD_DIM_VO = 128; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto USE_SLIDING_WINDOW = false; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + // The QKV matrices. + DTypeQ* q_ptr; + DTypeKV* k_ptr; + DTypeKV* v_ptr; + DTypeO* o_ptr; + float* lse_ptr; + + // Additional params + struct AdditionalParams { + double logits_soft_cap; +double sm_scale; +; + } additional_params; + + int64_t q_stride_n; + int64_t k_stride_n; + int64_t v_stride_n; + int64_t o_stride_n; + int64_t q_stride_h; + int64_t k_stride_h; + int64_t v_stride_h; + int64_t o_stride_h; + + int qo_len; + int kv_len; + int head_dim; + int num_qo_heads; + int num_kv_heads; + int group_size; + int window_left; + + bool causal; +}; + +#include \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_jit_pybind.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..2d6dd6076239658b22af333487f1f0fb49a86b1b --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_jit_pybind.cu @@ -0,0 +1,27 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "pytorch_extension_utils.h" +#include "single_prefill_sm90_config.inc" + +void single_prefill_with_kv_cache_sm90(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Single-request prefill attention with KV-Cache operator + m.def("run", single_prefill_with_kv_cache_sm90); +} diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_0.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..1f0930fe51f67c4b1a805a7301222d648998bb0c --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_0.cu @@ -0,0 +1,12 @@ +#include +#include "single_prefill_sm90_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SinglePrefillWithKVCacheDispatched + <128, 128, MaskMode::kNone, /*USE_SLIDING_WINDOW=*/false, DefaultAttention, Params>( + Params& params, cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_1.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..491c13341894b48e93c5e2da9d2a66160ac7bdfe --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_1.cu @@ -0,0 +1,12 @@ +#include +#include "single_prefill_sm90_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SinglePrefillWithKVCacheDispatched + <128, 128, MaskMode::kCausal, /*USE_SLIDING_WINDOW=*/false, DefaultAttention, Params>( + Params& params, cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_2.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..542f3f390d1514cbee2583533d9c269140d340ed --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_2.cu @@ -0,0 +1,12 @@ +#include +#include "single_prefill_sm90_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SinglePrefillWithKVCacheDispatched + <128, 128, MaskMode::kCustom, /*USE_SLIDING_WINDOW=*/false, DefaultAttention, Params>( + Params& params, cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_3.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..981f41acd163d7ef539b7750bd1ec43855e68db4 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_128_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_3.cu @@ -0,0 +1,12 @@ +#include +#include "single_prefill_sm90_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SinglePrefillWithKVCacheDispatched + <128, 128, MaskMode::kMultiItemScoring, /*USE_SLIDING_WINDOW=*/false, DefaultAttention, Params>( + Params& params, cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90.cu new file mode 100644 index 0000000000000000000000000000000000000000..39cabb9c2873d8e3922f6618e7081df62cac1b57 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90.cu @@ -0,0 +1,93 @@ +/* + * Copyright (c) 2024 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include +#include + +#include "pytorch_extension_utils.h" +#include "single_prefill_sm90_config.inc" + +namespace flashinfer { + +template +cudaError_t SinglePrefillWithKVCacheDispatched(Params& params, cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +void single_prefill_with_kv_cache_sm90(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS) { + unsigned int head_dim_qk = q.size(2); + unsigned int head_dim_vo = v.size(2); + unsigned int num_qo_heads = q.size(1); + unsigned int qo_len = q.size(0); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + QKVLayout kv_layout = static_cast(layout); + const c10::cuda::OptionalCUDAGuard device_guard(q.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + const MaskMode mask_mode = static_cast(mask_mode_code); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, USE_SLIDING_WINDOW, + USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + Params params; + params.q_ptr = static_cast(q.data_ptr()); + params.k_ptr = static_cast(k.data_ptr()); + params.v_ptr = static_cast(v.data_ptr()); + params.o_ptr = static_cast(o.data_ptr()); + params.lse_ptr = maybe_lse ? (static_cast(maybe_lse->data_ptr())) : nullptr; + params.q_stride_n = q.stride(0); + params.q_stride_h = q.stride(1); + params.o_stride_n = o.stride(0); + params.o_stride_h = o.stride(1); + if (kv_layout == QKVLayout::kNHD) { + params.k_stride_n = k.stride(0); + params.k_stride_h = k.stride(1); + params.v_stride_n = v.stride(0); + params.v_stride_h = v.stride(1); + } else { + params.k_stride_h = k.stride(0); + params.k_stride_n = k.stride(1); + params.v_stride_h = v.stride(0); + params.v_stride_n = v.stride(1); + } + params.qo_len = q.size(0); + params.kv_len = k.size(0); + params.num_qo_heads = q.size(1); + params.num_kv_heads = k.size(1); + params.causal = mask_mode == MaskMode::kCausal; + params.group_size = params.num_qo_heads / params.num_kv_heads; + params.window_left = window_left; + + ADDITIONAL_PARAMS_SETTER + + cudaError_t status = + SinglePrefillWithKVCacheDispatched(params, + stream); + TORCH_CHECK(status == cudaSuccess, "single_prefill_with_kv_cache_sm90 failed with error: " + + std::string(cudaGetErrorString(status))); + return true; + }); +} diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_config.inc b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..e9dd2831e4697053e92662f1d199c0e8a4b680c4 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_config.inc @@ -0,0 +1,70 @@ +#pragma once +#include +#include +#include +#include +#include +#include +#include + + +#define ADDITIONAL_FUNC_PARAMS , double logits_soft_cap, double sm_scale +#define ADDITIONAL_PARAMS_SETTER params.additional_params.logits_soft_cap = logits_soft_cap; \ +params.additional_params.sm_scale = sm_scale; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { using AttentionVariant = DefaultAttention; __VA_ARGS__(); }) + +using namespace flashinfer; + +using DTypeQ = cutlass_dtype_t; +using DTypeKV = cutlass_dtype_t; +using DTypeO = cutlass_dtype_t; +using IdType = cutlass_dtype_t; + +constexpr int HEAD_DIM_QK = 192; +constexpr int HEAD_DIM_VO = 128; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto USE_SLIDING_WINDOW = false; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + // The QKV matrices. + DTypeQ* q_ptr; + DTypeKV* k_ptr; + DTypeKV* v_ptr; + DTypeO* o_ptr; + float* lse_ptr; + + // Additional params + struct AdditionalParams { + double logits_soft_cap; +double sm_scale; +; + } additional_params; + + int64_t q_stride_n; + int64_t k_stride_n; + int64_t v_stride_n; + int64_t o_stride_n; + int64_t q_stride_h; + int64_t k_stride_h; + int64_t v_stride_h; + int64_t o_stride_h; + + int qo_len; + int kv_len; + int head_dim; + int num_qo_heads; + int num_kv_heads; + int group_size; + int window_left; + + bool causal; +}; + +#include \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_jit_pybind.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..2d6dd6076239658b22af333487f1f0fb49a86b1b --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_jit_pybind.cu @@ -0,0 +1,27 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "pytorch_extension_utils.h" +#include "single_prefill_sm90_config.inc" + +void single_prefill_with_kv_cache_sm90(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Single-request prefill attention with KV-Cache operator + m.def("run", single_prefill_with_kv_cache_sm90); +} diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_0.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..1a1584eebbf276192085ee7fda13999d2e07580f --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_0.cu @@ -0,0 +1,12 @@ +#include +#include "single_prefill_sm90_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SinglePrefillWithKVCacheDispatched + <192, 128, MaskMode::kNone, /*USE_SLIDING_WINDOW=*/false, DefaultAttention, Params>( + Params& params, cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_1.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..a0b314ab23075044a3f5367515796174e6a3391e --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_1.cu @@ -0,0 +1,12 @@ +#include +#include "single_prefill_sm90_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SinglePrefillWithKVCacheDispatched + <192, 128, MaskMode::kCausal, /*USE_SLIDING_WINDOW=*/false, DefaultAttention, Params>( + Params& params, cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_2.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..020022cd19ca3aaa434424349beede0edf06bbac --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_2.cu @@ -0,0 +1,12 @@ +#include +#include "single_prefill_sm90_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SinglePrefillWithKVCacheDispatched + <192, 128, MaskMode::kCustom, /*USE_SLIDING_WINDOW=*/false, DefaultAttention, Params>( + Params& params, cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_3.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..7eb0d82020eec67bc89f0cabf40c492aa055e3f5 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_192_head_dim_vo_128_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False_sm90/single_prefill_sm90_kernel_mask_3.cu @@ -0,0 +1,12 @@ +#include +#include "single_prefill_sm90_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SinglePrefillWithKVCacheDispatched + <192, 128, MaskMode::kMultiItemScoring, /*USE_SLIDING_WINDOW=*/false, DefaultAttention, Params>( + Params& params, cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill.cu new file mode 100644 index 0000000000000000000000000000000000000000..8d77f8420680e70a7929d8b16bd95f5587a01752 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill.cu @@ -0,0 +1,114 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include + +#include "flashinfer/fastdiv.cuh" +#include "pytorch_extension_utils.h" +#include "single_prefill_config.inc" + +namespace flashinfer { + +template +cudaError_t SinglePrefillWithKVCacheDispatched(Params params, typename Params::DTypeO* tmp, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +void single_prefill_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS) { + auto device = q.device(); + unsigned int head_dim_qk = q.size(2); + unsigned int kv_len, qo_len, num_kv_heads, num_qo_heads; + QKVLayout kv_layout = static_cast(layout); + qo_len = q.size(0); + num_qo_heads = q.size(1); + uint32_t q_stride_n = q.stride(0), q_stride_h = q.stride(1), k_stride_n, k_stride_h, v_stride_n, + v_stride_h; + if (kv_layout == QKVLayout::kNHD) { + kv_len = k.size(0); + num_kv_heads = k.size(1); + k_stride_n = k.stride(0); + k_stride_h = k.stride(1); + v_stride_n = v.stride(0); + v_stride_h = v.stride(1); + } else { + kv_len = k.size(1); + num_kv_heads = k.size(0); + k_stride_h = k.stride(0); + k_stride_n = k.stride(1); + v_stride_h = v.stride(0); + v_stride_n = v.stride(1); + } + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == qo_len, lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == num_qo_heads, lse.size(1), q.size(1)); + } + + const MaskMode mask_mode = static_cast(mask_mode_code); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(device); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, Params, + [&] { + Params params; + + params.q = static_cast(q.data_ptr()); + params.k = static_cast(k.data_ptr()); + params.v = static_cast(v.data_ptr()); + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.num_qo_heads = num_qo_heads; + params.num_kv_heads = num_kv_heads; + params.group_size = uint_fastdiv(num_qo_heads / num_kv_heads); + params.qo_len = qo_len; + params.kv_len = kv_len; + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.k_stride_n = k_stride_n; + params.k_stride_h = k_stride_h; + params.v_stride_n = v_stride_n; + params.v_stride_h = v_stride_h; + + params.window_left = window_left; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + cudaError_t status = flashinfer::SinglePrefillWithKVCacheDispatched< + HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + /*use_fp16_qk_reduction=*/USE_FP16_QK_REDUCTION, MASK_MODE, AttentionVariant>( + params, static_cast(tmp.data_ptr()), stream); + TORCH_CHECK(status == cudaSuccess, + "SinglePrefillWithKVCache kernel launch failed, error: " + + std::string(cudaGetErrorString(status))); + return true; + }); +} diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_config.inc b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..80ac4f4cb1312daf52e7416a9dfe5ac809b53754 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_config.inc @@ -0,0 +1,83 @@ +#pragma once +#include +#include +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_custom_mask, std::optional maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta +#define ADDITIONAL_PARAMS_SETTER params.maybe_custom_mask = maybe_custom_mask ? static_cast(maybe_custom_mask->data_ptr()): nullptr; \ +params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; + + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, Params, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { \ + constexpr bool use_custom_mask = MASK_MODE == MaskMode::kCustom; \ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ + }) + + +using namespace flashinfer; + +using DTypeQ = half; +using DTypeKV = half; +using DTypeO = half; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 256; +constexpr int HEAD_DIM_VO = 256; +constexpr bool USE_FP16_QK_REDUCTION = false; +constexpr auto USE_LOGITS_SOFT_CAP = true; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = true; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = int32_t; + DTypeQ* q; + DTypeKV* k; + DTypeKV* v; + DTypeO* o; + float* lse; + uint_fastdiv group_size; + + uint8_t* maybe_custom_mask; +float* maybe_alibi_slopes; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; + + + uint32_t qo_len; + uint32_t kv_len; + uint32_t num_qo_heads; + uint32_t num_kv_heads; + uint32_t q_stride_n; + uint32_t q_stride_h; + uint32_t k_stride_n; + uint32_t k_stride_h; + uint32_t v_stride_n; + uint32_t v_stride_h; + uint32_t head_dim; + int32_t window_left; + + bool partition_kv; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { + return qo_len; + } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return kv_len; + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_jit_pybind.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..aaccbc045f9fd8566654a106379c99a053489fc7 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_jit_pybind.cu @@ -0,0 +1,27 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "pytorch_extension_utils.h" +#include "single_prefill_config.inc" + +void single_prefill_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Single-request prefill attention with KV-Cache operator + m.def("run", single_prefill_with_kv_cache); +} diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_kernel_mask_0.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..14843c3fb10f36c123129d10410958397c0add42 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_kernel_mask_0.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kNone == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 256, 256, PosEncodingMode::kNone, false, MaskMode::kNone, DefaultAttention, Params>( + Params params, half* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_kernel_mask_1.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..dd1c37887e9c496433beabd47a6c0e627a6a5468 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_kernel_mask_1.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCausal == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 256, 256, PosEncodingMode::kNone, false, MaskMode::kCausal, DefaultAttention, Params>( + Params params, half* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_kernel_mask_2.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..486980ebafeeba7400fc788925911ccb7337400e --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_kernel_mask_2.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCustom == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 256, 256, PosEncodingMode::kNone, false, MaskMode::kCustom, DefaultAttention, Params>( + Params params, half* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_kernel_mask_3.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..a4c568fac4ff8089c8e50ec27b322291cabf84da --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False/single_prefill_kernel_mask_3.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kMultiItemScoring == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 256, 256, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, DefaultAttention, Params>( + Params params, half* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90.cu new file mode 100644 index 0000000000000000000000000000000000000000..39cabb9c2873d8e3922f6618e7081df62cac1b57 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90.cu @@ -0,0 +1,93 @@ +/* + * Copyright (c) 2024 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include +#include + +#include "pytorch_extension_utils.h" +#include "single_prefill_sm90_config.inc" + +namespace flashinfer { + +template +cudaError_t SinglePrefillWithKVCacheDispatched(Params& params, cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +void single_prefill_with_kv_cache_sm90(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS) { + unsigned int head_dim_qk = q.size(2); + unsigned int head_dim_vo = v.size(2); + unsigned int num_qo_heads = q.size(1); + unsigned int qo_len = q.size(0); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + QKVLayout kv_layout = static_cast(layout); + const c10::cuda::OptionalCUDAGuard device_guard(q.device()); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + const MaskMode mask_mode = static_cast(mask_mode_code); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, USE_SLIDING_WINDOW, + USE_LOGITS_SOFT_CAP, AttentionVariant, Params, [&] { + Params params; + params.q_ptr = static_cast(q.data_ptr()); + params.k_ptr = static_cast(k.data_ptr()); + params.v_ptr = static_cast(v.data_ptr()); + params.o_ptr = static_cast(o.data_ptr()); + params.lse_ptr = maybe_lse ? (static_cast(maybe_lse->data_ptr())) : nullptr; + params.q_stride_n = q.stride(0); + params.q_stride_h = q.stride(1); + params.o_stride_n = o.stride(0); + params.o_stride_h = o.stride(1); + if (kv_layout == QKVLayout::kNHD) { + params.k_stride_n = k.stride(0); + params.k_stride_h = k.stride(1); + params.v_stride_n = v.stride(0); + params.v_stride_h = v.stride(1); + } else { + params.k_stride_h = k.stride(0); + params.k_stride_n = k.stride(1); + params.v_stride_h = v.stride(0); + params.v_stride_n = v.stride(1); + } + params.qo_len = q.size(0); + params.kv_len = k.size(0); + params.num_qo_heads = q.size(1); + params.num_kv_heads = k.size(1); + params.causal = mask_mode == MaskMode::kCausal; + params.group_size = params.num_qo_heads / params.num_kv_heads; + params.window_left = window_left; + + ADDITIONAL_PARAMS_SETTER + + cudaError_t status = + SinglePrefillWithKVCacheDispatched(params, + stream); + TORCH_CHECK(status == cudaSuccess, "single_prefill_with_kv_cache_sm90 failed with error: " + + std::string(cudaGetErrorString(status))); + return true; + }); +} diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_config.inc b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..63f8975cb37415bb018066284f37bed6e60442f4 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_config.inc @@ -0,0 +1,70 @@ +#pragma once +#include +#include +#include +#include +#include +#include +#include + + +#define ADDITIONAL_FUNC_PARAMS , double logits_soft_cap, double sm_scale +#define ADDITIONAL_PARAMS_SETTER params.additional_params.logits_soft_cap = logits_soft_cap; \ +params.additional_params.sm_scale = sm_scale; + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { using AttentionVariant = DefaultAttention; __VA_ARGS__(); }) + +using namespace flashinfer; + +using DTypeQ = cutlass_dtype_t; +using DTypeKV = cutlass_dtype_t; +using DTypeO = cutlass_dtype_t; +using IdType = cutlass_dtype_t; + +constexpr int HEAD_DIM_QK = 256; +constexpr int HEAD_DIM_VO = 256; +constexpr auto USE_LOGITS_SOFT_CAP = true; +constexpr auto USE_SLIDING_WINDOW = true; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = IdType; + + // The QKV matrices. + DTypeQ* q_ptr; + DTypeKV* k_ptr; + DTypeKV* v_ptr; + DTypeO* o_ptr; + float* lse_ptr; + + // Additional params + struct AdditionalParams { + double logits_soft_cap; +double sm_scale; +; + } additional_params; + + int64_t q_stride_n; + int64_t k_stride_n; + int64_t v_stride_n; + int64_t o_stride_n; + int64_t q_stride_h; + int64_t k_stride_h; + int64_t v_stride_h; + int64_t o_stride_h; + + int qo_len; + int kv_len; + int head_dim; + int num_qo_heads; + int num_kv_heads; + int group_size; + int window_left; + + bool causal; +}; + +#include \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_jit_pybind.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..2d6dd6076239658b22af333487f1f0fb49a86b1b --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_jit_pybind.cu @@ -0,0 +1,27 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "pytorch_extension_utils.h" +#include "single_prefill_sm90_config.inc" + +void single_prefill_with_kv_cache_sm90(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Single-request prefill attention with KV-Cache operator + m.def("run", single_prefill_with_kv_cache_sm90); +} diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_kernel_mask_0.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..5b881dbe8c829acac6560accc4ea0a629730ae05 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_kernel_mask_0.cu @@ -0,0 +1,12 @@ +#include +#include "single_prefill_sm90_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SinglePrefillWithKVCacheDispatched + <256, 256, MaskMode::kNone, /*USE_SLIDING_WINDOW=*/true, DefaultAttention, Params>( + Params& params, cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_kernel_mask_1.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..e081f468769e66fd45aef98e65f0c93f5f667745 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_kernel_mask_1.cu @@ -0,0 +1,12 @@ +#include +#include "single_prefill_sm90_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SinglePrefillWithKVCacheDispatched + <256, 256, MaskMode::kCausal, /*USE_SLIDING_WINDOW=*/true, DefaultAttention, Params>( + Params& params, cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_kernel_mask_2.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..ac76c050bef07556fb25a86da5c1c6148d856db5 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_kernel_mask_2.cu @@ -0,0 +1,12 @@ +#include +#include "single_prefill_sm90_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SinglePrefillWithKVCacheDispatched + <256, 256, MaskMode::kCustom, /*USE_SLIDING_WINDOW=*/true, DefaultAttention, Params>( + Params& params, cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_kernel_mask_3.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..6b5e45f1881bea91f4154b5d0968cc5a5750176d --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_256_head_dim_vo_256_posenc_0_use_swa_True_use_logits_cap_True_f16qk_False_sm90/single_prefill_sm90_kernel_mask_3.cu @@ -0,0 +1,12 @@ +#include +#include "single_prefill_sm90_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +template cudaError_t SinglePrefillWithKVCacheDispatched + <256, 256, MaskMode::kMultiItemScoring, /*USE_SLIDING_WINDOW=*/true, DefaultAttention, Params>( + Params& params, cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill.cu new file mode 100644 index 0000000000000000000000000000000000000000..8d77f8420680e70a7929d8b16bd95f5587a01752 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill.cu @@ -0,0 +1,114 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include + +#include "flashinfer/fastdiv.cuh" +#include "pytorch_extension_utils.h" +#include "single_prefill_config.inc" + +namespace flashinfer { + +template +cudaError_t SinglePrefillWithKVCacheDispatched(Params params, typename Params::DTypeO* tmp, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +void single_prefill_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS) { + auto device = q.device(); + unsigned int head_dim_qk = q.size(2); + unsigned int kv_len, qo_len, num_kv_heads, num_qo_heads; + QKVLayout kv_layout = static_cast(layout); + qo_len = q.size(0); + num_qo_heads = q.size(1); + uint32_t q_stride_n = q.stride(0), q_stride_h = q.stride(1), k_stride_n, k_stride_h, v_stride_n, + v_stride_h; + if (kv_layout == QKVLayout::kNHD) { + kv_len = k.size(0); + num_kv_heads = k.size(1); + k_stride_n = k.stride(0); + k_stride_h = k.stride(1); + v_stride_n = v.stride(0); + v_stride_h = v.stride(1); + } else { + kv_len = k.size(1); + num_kv_heads = k.size(0); + k_stride_h = k.stride(0); + k_stride_n = k.stride(1); + v_stride_h = v.stride(0); + v_stride_n = v.stride(1); + } + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == qo_len, lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == num_qo_heads, lse.size(1), q.size(1)); + } + + const MaskMode mask_mode = static_cast(mask_mode_code); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(device); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, Params, + [&] { + Params params; + + params.q = static_cast(q.data_ptr()); + params.k = static_cast(k.data_ptr()); + params.v = static_cast(v.data_ptr()); + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.num_qo_heads = num_qo_heads; + params.num_kv_heads = num_kv_heads; + params.group_size = uint_fastdiv(num_qo_heads / num_kv_heads); + params.qo_len = qo_len; + params.kv_len = kv_len; + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.k_stride_n = k_stride_n; + params.k_stride_h = k_stride_h; + params.v_stride_n = v_stride_n; + params.v_stride_h = v_stride_h; + + params.window_left = window_left; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + cudaError_t status = flashinfer::SinglePrefillWithKVCacheDispatched< + HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + /*use_fp16_qk_reduction=*/USE_FP16_QK_REDUCTION, MASK_MODE, AttentionVariant>( + params, static_cast(tmp.data_ptr()), stream); + TORCH_CHECK(status == cudaSuccess, + "SinglePrefillWithKVCache kernel launch failed, error: " + + std::string(cudaGetErrorString(status))); + return true; + }); +} diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_config.inc b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..3ee379d4da5f2f8dc5379111173c7938a8ecaddf --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_config.inc @@ -0,0 +1,83 @@ +#pragma once +#include +#include +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_custom_mask, std::optional maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta +#define ADDITIONAL_PARAMS_SETTER params.maybe_custom_mask = maybe_custom_mask ? static_cast(maybe_custom_mask->data_ptr()): nullptr; \ +params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; + + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, Params, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { \ + constexpr bool use_custom_mask = MASK_MODE == MaskMode::kCustom; \ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ + }) + + +using namespace flashinfer; + +using DTypeQ = half; +using DTypeKV = half; +using DTypeO = half; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 64; +constexpr int HEAD_DIM_VO = 64; +constexpr bool USE_FP16_QK_REDUCTION = false; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = false; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = int32_t; + DTypeQ* q; + DTypeKV* k; + DTypeKV* v; + DTypeO* o; + float* lse; + uint_fastdiv group_size; + + uint8_t* maybe_custom_mask; +float* maybe_alibi_slopes; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; + + + uint32_t qo_len; + uint32_t kv_len; + uint32_t num_qo_heads; + uint32_t num_kv_heads; + uint32_t q_stride_n; + uint32_t q_stride_h; + uint32_t k_stride_n; + uint32_t k_stride_h; + uint32_t v_stride_n; + uint32_t v_stride_h; + uint32_t head_dim; + int32_t window_left; + + bool partition_kv; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { + return qo_len; + } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return kv_len; + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_jit_pybind.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..aaccbc045f9fd8566654a106379c99a053489fc7 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_jit_pybind.cu @@ -0,0 +1,27 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "pytorch_extension_utils.h" +#include "single_prefill_config.inc" + +void single_prefill_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Single-request prefill attention with KV-Cache operator + m.def("run", single_prefill_with_kv_cache); +} diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_0.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..465ba8aa40f1852772e15b84eace44711f34a177 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_0.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kNone == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, DefaultAttention, Params>( + Params params, half* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_1.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..1fd3d0a416bf7cfacd861519ac03ed387b5b6526 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_1.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCausal == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, DefaultAttention, Params>( + Params params, half* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_2.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..f71cb040b06f03780749c54b00cbae78a3872b68 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_2.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCustom == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, DefaultAttention, Params>( + Params params, half* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_3.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..6800c5a9275289669e7107ded5bb7bc13245b5f6 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_False_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_3.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kMultiItemScoring == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, DefaultAttention, Params>( + Params params, half* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill.cu new file mode 100644 index 0000000000000000000000000000000000000000..8d77f8420680e70a7929d8b16bd95f5587a01752 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill.cu @@ -0,0 +1,114 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include + +#include "flashinfer/fastdiv.cuh" +#include "pytorch_extension_utils.h" +#include "single_prefill_config.inc" + +namespace flashinfer { + +template +cudaError_t SinglePrefillWithKVCacheDispatched(Params params, typename Params::DTypeO* tmp, + cudaStream_t stream); + +} // namespace flashinfer + +using namespace flashinfer; + +void single_prefill_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS) { + auto device = q.device(); + unsigned int head_dim_qk = q.size(2); + unsigned int kv_len, qo_len, num_kv_heads, num_qo_heads; + QKVLayout kv_layout = static_cast(layout); + qo_len = q.size(0); + num_qo_heads = q.size(1); + uint32_t q_stride_n = q.stride(0), q_stride_h = q.stride(1), k_stride_n, k_stride_h, v_stride_n, + v_stride_h; + if (kv_layout == QKVLayout::kNHD) { + kv_len = k.size(0); + num_kv_heads = k.size(1); + k_stride_n = k.stride(0); + k_stride_h = k.stride(1); + v_stride_n = v.stride(0); + v_stride_h = v.stride(1); + } else { + kv_len = k.size(1); + num_kv_heads = k.size(0); + k_stride_h = k.stride(0); + k_stride_n = k.stride(1); + v_stride_h = v.stride(0); + v_stride_n = v.stride(1); + } + if (maybe_lse) { + const auto& lse = *maybe_lse; + TORCH_CHECK(lse.size(0) == qo_len, lse.size(0), q.size(0)); + TORCH_CHECK(lse.size(1) == num_qo_heads, lse.size(1), q.size(1)); + } + + const MaskMode mask_mode = static_cast(mask_mode_code); + + auto q_scalar_type = q.scalar_type(); + auto kv_scalar_type = k.scalar_type(); + + const c10::cuda::OptionalCUDAGuard device_guard(device); + const cudaStream_t stream = c10::cuda::getCurrentCUDAStream(); + + DISPATCH_context( + DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, Params, + [&] { + Params params; + + params.q = static_cast(q.data_ptr()); + params.k = static_cast(k.data_ptr()); + params.v = static_cast(v.data_ptr()); + params.o = static_cast(o.data_ptr()); + params.lse = maybe_lse ? static_cast(maybe_lse->data_ptr()) : nullptr; + params.num_qo_heads = num_qo_heads; + params.num_kv_heads = num_kv_heads; + params.group_size = uint_fastdiv(num_qo_heads / num_kv_heads); + params.qo_len = qo_len; + params.kv_len = kv_len; + params.q_stride_n = q_stride_n; + params.q_stride_h = q_stride_h; + params.k_stride_n = k_stride_n; + params.k_stride_h = k_stride_h; + params.v_stride_n = v_stride_n; + params.v_stride_h = v_stride_h; + + params.window_left = window_left; + params.partition_kv = false; + + ADDITIONAL_PARAMS_SETTER + + cudaError_t status = flashinfer::SinglePrefillWithKVCacheDispatched< + HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, + /*use_fp16_qk_reduction=*/USE_FP16_QK_REDUCTION, MASK_MODE, AttentionVariant>( + params, static_cast(tmp.data_ptr()), stream); + TORCH_CHECK(status == cudaSuccess, + "SinglePrefillWithKVCache kernel launch failed, error: " + + std::string(cudaGetErrorString(status))); + return true; + }); +} diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_config.inc b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_config.inc new file mode 100644 index 0000000000000000000000000000000000000000..bcb29930dfc32590640c628f8ac400bfccacb34c --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_config.inc @@ -0,0 +1,83 @@ +#pragma once +#include +#include +#include +#include +#include +#include + +#define ADDITIONAL_FUNC_PARAMS , std::optional maybe_custom_mask, std::optional maybe_alibi_slopes, double logits_soft_cap, double sm_scale, double rope_rcp_scale, double rope_rcp_theta +#define ADDITIONAL_PARAMS_SETTER params.maybe_custom_mask = maybe_custom_mask ? static_cast(maybe_custom_mask->data_ptr()): nullptr; \ +params.maybe_alibi_slopes = maybe_alibi_slopes ? static_cast(maybe_alibi_slopes->data_ptr()): nullptr; \ +params.logits_soft_cap = logits_soft_cap; \ +params.sm_scale = sm_scale; \ +params.rope_rcp_scale = rope_rcp_scale; \ +params.rope_rcp_theta = rope_rcp_theta; + + +#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, Params, ...) \ + DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { \ + constexpr bool use_custom_mask = MASK_MODE == MaskMode::kCustom; \ + using AttentionVariant = DefaultAttention; \ + __VA_ARGS__(); \ + }) + + +using namespace flashinfer; + +using DTypeQ = half; +using DTypeKV = half; +using DTypeO = half; +using IdType = int32_t; +constexpr int HEAD_DIM_QK = 64; +constexpr int HEAD_DIM_VO = 64; +constexpr bool USE_FP16_QK_REDUCTION = false; +constexpr auto USE_LOGITS_SOFT_CAP = false; +constexpr auto POS_ENCODING_MODE = PosEncodingMode::kNone; +constexpr auto USE_SLIDING_WINDOW = true; + +struct Params { + using DTypeQ = DTypeQ; + using DTypeKV = DTypeKV; + using DTypeO = DTypeO; + using IdType = int32_t; + DTypeQ* q; + DTypeKV* k; + DTypeKV* v; + DTypeO* o; + float* lse; + uint_fastdiv group_size; + + uint8_t* maybe_custom_mask; +float* maybe_alibi_slopes; +double logits_soft_cap; +double sm_scale; +double rope_rcp_scale; +double rope_rcp_theta; + + + uint32_t qo_len; + uint32_t kv_len; + uint32_t num_qo_heads; + uint32_t num_kv_heads; + uint32_t q_stride_n; + uint32_t q_stride_h; + uint32_t k_stride_n; + uint32_t k_stride_h; + uint32_t v_stride_n; + uint32_t v_stride_h; + uint32_t head_dim; + int32_t window_left; + + bool partition_kv; + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { + return qo_len; + } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return kv_len; + } +}; + +#include \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_jit_pybind.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_jit_pybind.cu new file mode 100644 index 0000000000000000000000000000000000000000..aaccbc045f9fd8566654a106379c99a053489fc7 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_jit_pybind.cu @@ -0,0 +1,27 @@ +/* + * Copyright (c) 2023-2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "pytorch_extension_utils.h" +#include "single_prefill_config.inc" + +void single_prefill_with_kv_cache(at::Tensor q, at::Tensor k, at::Tensor v, at::Tensor tmp, + at::Tensor o, std::optional maybe_lse, + int64_t mask_mode_code, int64_t layout, + int64_t window_left ADDITIONAL_FUNC_PARAMS); + +TORCH_LIBRARY_FRAGMENT(TORCH_EXTENSION_NAME, m) { + // Single-request prefill attention with KV-Cache operator + m.def("run", single_prefill_with_kv_cache); +} diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_0.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_0.cu new file mode 100644 index 0000000000000000000000000000000000000000..105f200cc3715c7f935961ea39f055f9f7a12b40 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_0.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kNone == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 64, 64, PosEncodingMode::kNone, false, MaskMode::kNone, DefaultAttention, Params>( + Params params, half* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_1.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_1.cu new file mode 100644 index 0000000000000000000000000000000000000000..c835246808e83388370fd3b22b6703a5db464580 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_1.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCausal == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 64, 64, PosEncodingMode::kNone, false, MaskMode::kCausal, DefaultAttention, Params>( + Params params, half* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_2.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_2.cu new file mode 100644 index 0000000000000000000000000000000000000000..8104e776733a3ed680359f01197b9a172ac43248 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_2.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kCustom == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 64, 64, PosEncodingMode::kNone, false, MaskMode::kCustom, DefaultAttention, Params>( + Params params, half* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_3.cu b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_3.cu new file mode 100644 index 0000000000000000000000000000000000000000..b30f6e965634630a057766e440ffeb2ebe9cebd9 --- /dev/null +++ b/csrc/generated/single_prefill_with_kv_cache_dtype_q_f16_dtype_kv_f16_dtype_o_f16_head_dim_qk_64_head_dim_vo_64_posenc_0_use_swa_True_use_logits_cap_False_f16qk_False/single_prefill_kernel_mask_3.cu @@ -0,0 +1,15 @@ +#include +#include "single_prefill_config.inc" + +using namespace flashinfer; + +namespace flashinfer { + +constexpr auto use_custom_mask = MaskMode::kMultiItemScoring == MaskMode::kCustom; + +template cudaError_t SinglePrefillWithKVCacheDispatched< + 64, 64, PosEncodingMode::kNone, false, MaskMode::kMultiItemScoring, DefaultAttention, Params>( + Params params, half* tmp, + cudaStream_t stream); + +}; \ No newline at end of file diff --git a/flake.lock b/flake.lock new file mode 100644 index 0000000000000000000000000000000000000000..612fbf05e4665569fb0186ab0b9e5d4e1551e9e0 --- /dev/null +++ b/flake.lock @@ -0,0 +1,168 @@ +{ + "nodes": { + "flake-compat": { + "locked": { + "lastModified": 1747046372, + "narHash": "sha256-CIVLLkVgvHYbgI2UpXvIIBJ12HWgX+fjA8Xf8PUmqCY=", + "owner": "edolstra", + "repo": "flake-compat", + "rev": "9100a0f413b0c601e0533d1d94ffd501ce2e7885", + "type": "github" + }, + "original": { + "owner": "edolstra", + "repo": "flake-compat", + "type": "github" + } + }, + "flake-compat_2": { + "locked": { + "lastModified": 1733328505, + "narHash": "sha256-NeCCThCEP3eCl2l/+27kNNK7QrwZB1IJCrXfrbv5oqU=", + "owner": "edolstra", + "repo": "flake-compat", + "rev": "ff81ac966bb2cae68946d5ed5fc4994f96d0ffec", + "type": "github" + }, + "original": { + "owner": "edolstra", + "repo": "flake-compat", + "type": "github" + } + }, + "flake-utils": { + "inputs": { + "systems": "systems" + }, + "locked": { + "lastModified": 1731533236, + "narHash": "sha256-l0KFg5HjrsfsO/JpG+r7fRrqm12kzFHyUHqHCVpMMbI=", + "owner": "numtide", + "repo": "flake-utils", + "rev": "11707dc2f618dd54ca8739b309ec4fc024de578b", + "type": "github" + }, + "original": { + "owner": "numtide", + "repo": "flake-utils", + "type": "github" + } + }, + "flake-utils_2": { + "inputs": { + "systems": "systems_2" + }, + "locked": { + "lastModified": 1731533236, + "narHash": "sha256-l0KFg5HjrsfsO/JpG+r7fRrqm12kzFHyUHqHCVpMMbI=", + "owner": "numtide", + "repo": "flake-utils", + "rev": "11707dc2f618dd54ca8739b309ec4fc024de578b", + "type": "github" + }, + "original": { + "owner": "numtide", + "repo": "flake-utils", + "type": "github" + } + }, + "hf-nix": { + "inputs": { + "flake-compat": "flake-compat_2", + "flake-utils": "flake-utils_2", + "nixpkgs": "nixpkgs" + }, + "locked": { + "lastModified": 1754038838, + "narHash": "sha256-oHigCT4z0ayyLyEuxdZooSXRAZP8lfOkZHzY1lx1U50=", + "owner": "huggingface", + "repo": "hf-nix", + "rev": "336f781fa284e193baa3d4c3ce3f95fb34e9ffad", + "type": "github" + }, + "original": { + "owner": "huggingface", + "repo": "hf-nix", + "type": "github" + } + }, + "kernel-builder": { + "inputs": { + "flake-compat": "flake-compat", + "flake-utils": "flake-utils", + "hf-nix": "hf-nix", + "nixpkgs": [ + "kernel-builder", + "hf-nix", + "nixpkgs" + ] + }, + "locked": { + "lastModified": 1754509142, + "narHash": "sha256-UKM847nSL/7WXvi7rBo0AuucI7wVTz9KGY4aZdDOjrI=", + "owner": "huggingface", + "repo": "kernel-builder", + "rev": "3b10038fed6416e1ab89c708b66207cfc0248e14", + "type": "github" + }, + "original": { + "owner": "huggingface", + "repo": "kernel-builder", + "type": "github" + } + }, + "nixpkgs": { + "locked": { + "lastModified": 1752785354, + "narHash": "sha256-Y33ryUz7MPqKrZwlbQcsYCUz2jAJCacRf8jbs0tYUlA=", + "owner": "nixos", + "repo": "nixpkgs", + "rev": "d38025438a6ee456758dc03188ca6873a415463b", + "type": "github" + }, + "original": { + "owner": "nixos", + "repo": "nixpkgs", + "rev": "d38025438a6ee456758dc03188ca6873a415463b", + "type": "github" + } + }, + "root": { + "inputs": { + "kernel-builder": "kernel-builder" + } + }, + "systems": { + "locked": { + "lastModified": 1681028828, + "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=", + "owner": "nix-systems", + "repo": "default", + "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e", + "type": "github" + }, + "original": { + "owner": "nix-systems", + "repo": "default", + "type": "github" + } + }, + "systems_2": { + "locked": { + "lastModified": 1681028828, + "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=", + "owner": "nix-systems", + "repo": "default", + "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e", + "type": "github" + }, + "original": { + "owner": "nix-systems", + "repo": "default", + "type": "github" + } + } + }, + "root": "root", + "version": 7 +} diff --git a/flake.nix b/flake.nix new file mode 100644 index 0000000000000000000000000000000000000000..a7e3d1c44c5c1baf5f28caf1afcda3691308616a --- /dev/null +++ b/flake.nix @@ -0,0 +1,17 @@ +{ + description = "Flake for flashinfer kernel"; + + inputs = { + kernel-builder.url = "github:huggingface/kernel-builder"; + }; + + outputs = + { + self, + kernel-builder, + }: + kernel-builder.lib.genFlakeOutputs { + path = ./.; + rev = self.shortRev or self.dirtyShortRev or self.lastModifiedDate; + }; +} diff --git a/generate-source.md b/generate-source.md new file mode 100644 index 0000000000000000000000000000000000000000..5ffba3c633895db05bc39917012666d83a792d03 --- /dev/null +++ b/generate-source.md @@ -0,0 +1,28 @@ +# Generate FlashInfer Source Files + +```bash +# 1. Clone FlashInfer (pin version with --branch) +git clone --depth 1 --branch v0.2.0 https://github.com/flashinfer-ai/flashinfer flashinfer + +# 2. Apply patch (if needed) +cd flashinfer && git apply ../generated.patch + +# 3. Generate AOT files +export CUDA_HOME=/usr/local/cuda-12.6 && export TORCH_CUDA_ARCH_LIST="7.5 8.0 8.6 8.7 8.9 9.0" && uv run --with torch python -m flashinfer.aot + +# 4. Copy generated files to csrc directory +cp -r build/aot/generated ../csrc/generated + +# 5. Push to HuggingFace +git remote add origin git@hf.co:kernels-community/flashinfer-draft +git push origin main +``` + +now manually comment out the `TORCH_LIBRARY_FRAGMENT` in the generated files + +specifically in the activation `csrc/generated/gelu_and_mul.cu` and etc since those are included in the current build. + +### build and test + +```bash +``` diff --git a/generated.patch b/generated.patch new file mode 100644 index 0000000000000000000000000000000000000000..d75b70b1602efb7ba44129f94a23b82056e51ab8 --- /dev/null +++ b/generated.patch @@ -0,0 +1,89 @@ +diff --git a/flashinfer/aot.py b/flashinfer/aot.py +index 2e2885a..981879f 100644 +--- a/flashinfer/aot.py ++++ b/flashinfer/aot.py +@@ -378,15 +378,29 @@ def gen_all_modules( + jit_specs.append(gen_vllm_comm_module()) + + if add_misc: +- jit_specs += [ ++ misc_modules = [ + gen_cascade_module(), + gen_norm_module(), +- gen_nvshmem_module(), ++ ] ++ ++ # Only add nvshmem module if nvidia.nvshmem is available or env paths are set ++ if (os.environ.get("NVSHMEM_INCLUDE_PATH") or ++ os.environ.get("NVSHMEM_LIBRARY_PATH")): ++ misc_modules.append(gen_nvshmem_module()) ++ else: ++ try: ++ import nvidia.nvshmem ++ misc_modules.append(gen_nvshmem_module()) ++ except ImportError: ++ pass # Skip nvshmem module if not available ++ ++ misc_modules += [ + gen_page_module(), + gen_quantization_module(), + gen_rope_module(), + gen_sampling_module(), + ] ++ jit_specs += misc_modules + if has_sm90: + jit_specs.append(get_trtllm_utils_spec()) + +@@ -657,12 +671,12 @@ def main(): + ) + print("Total ops:", len(jit_specs)) + +- # Build +- build_jit_specs(jit_specs, verbose=True, skip_prebuilt=False) ++ # # Build ++ # build_jit_specs(jit_specs, verbose=True, skip_prebuilt=False) + +- # Copy built kernels +- copy_built_kernels(jit_specs, out_dir) +- print("AOT kernels saved to:", out_dir) ++ # # Copy built kernels ++ # copy_built_kernels(jit_specs, out_dir) ++ # print("AOT kernels saved to:", out_dir) + + + if __name__ == "__main__": +diff --git a/flashinfer/jit/env.py b/flashinfer/jit/env.py +index 24104ee..5732a8f 100644 +--- a/flashinfer/jit/env.py ++++ b/flashinfer/jit/env.py +@@ -73,10 +73,12 @@ def get_nvshmem_include_dirs(): + if paths is not None: + return [pathlib.Path(p) for p in paths.split(os.pathsep) if p] + +- import nvidia.nvshmem +- +- path = pathlib.Path(nvidia.nvshmem.__path__[0]) / "include" +- return [path] ++ try: ++ import nvidia.nvshmem ++ path = pathlib.Path(nvidia.nvshmem.__path__[0]) / "include" ++ return [path] ++ except ImportError: ++ return [] + + + def get_nvshmem_lib_dirs(): +@@ -84,7 +86,9 @@ def get_nvshmem_lib_dirs(): + if paths is not None: + return [pathlib.Path(p) for p in paths.split(os.pathsep) if p] + +- import nvidia.nvshmem +- +- path = pathlib.Path(nvidia.nvshmem.__path__[0]) / "lib" +- return [path] ++ try: ++ import nvidia.nvshmem ++ path = pathlib.Path(nvidia.nvshmem.__path__[0]) / "lib" ++ return [path] ++ except ImportError: ++ return [] diff --git a/include/flashinfer/activation.cuh b/include/flashinfer/activation.cuh new file mode 100644 index 0000000000000000000000000000000000000000..6e9f02992335115d71b6d60c114ece2104a97a65 --- /dev/null +++ b/include/flashinfer/activation.cuh @@ -0,0 +1,69 @@ +/* + * Copyright (c) 2024 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#ifndef FLASHINFER_ACTIVATION_CUH_ +#define FLASHINFER_ACTIVATION_CUH_ + +#include "math.cuh" +#include "utils.cuh" +#include "vec_dtypes.cuh" + +namespace flashinfer { + +namespace activation { + +template +__global__ void act_and_mul_kernel(T* __restrict__ out, const T* __restrict__ input, const int d) { + constexpr uint32_t vec_size = 16 / sizeof(T); + const int64_t token_idx = blockIdx.x; + const int64_t thread_idx = threadIdx.x; + const int64_t stride = blockDim.x; + const int64_t offset = token_idx * 2 * d; + +#if (__CUDACC_VER_MAJOR__ >= 12 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.wait;"); +#endif + +#pragma unroll 1 + for (uint32_t idx = thread_idx; idx < d / vec_size; idx += stride) { + vec_t x_vec, y_vec, out_vec; + x_vec.cast_load(input + offset + idx * vec_size); + y_vec.cast_load(input + offset + d + idx * vec_size); +#pragma unroll + for (uint32_t i = 0; i < vec_size; ++i) { + out_vec[i] = Activation(x_vec[i]) * y_vec[i]; + } + out_vec.cast_store(out + token_idx * d + idx * vec_size); + } + + const int64_t remaining_offset = d - d % (stride * vec_size); + // process the remaining elements +#pragma unroll 1 + for (int64_t idx = thread_idx; idx < d % (stride * vec_size); idx += stride) { + float x = input[offset + remaining_offset + idx], + y = input[offset + remaining_offset + d + idx]; + out[token_idx * d + remaining_offset + idx] = Activation(x) * y; + } + +#if (__CUDACC_VER_MAJOR__ >= 12 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.launch_dependents;"); +#endif +} + +} // namespace activation +} // namespace flashinfer + +#endif // FLASHINFER_ACTIVATION_CUH_ diff --git a/include/flashinfer/allocator.h b/include/flashinfer/allocator.h new file mode 100644 index 0000000000000000000000000000000000000000..e4eb767cd75fb320f302d2ee4f99937bf9936710 --- /dev/null +++ b/include/flashinfer/allocator.h @@ -0,0 +1,63 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef FLASHINFER_ALLOCATOR_H_ +#define FLASHINFER_ALLOCATOR_H_ + +#include +#include + +#include "exception.h" + +namespace flashinfer { + +// create a function that returns T* from base pointer and offset +template +T* GetPtrFromBaseOffset(void* base_ptr, int64_t offset) { + return reinterpret_cast(reinterpret_cast(base_ptr) + offset); +} + +struct AlignedAllocator { + void* base_ptr; + void* cur_ptr; + size_t remaining_space; + AlignedAllocator(void* buf, size_t space) : base_ptr(buf), cur_ptr(buf), remaining_space(space) {} + template + T* aligned_alloc(size_t size, size_t alignment, std::string name) { + if (std::align(alignment, size, cur_ptr, remaining_space)) { + T* result = reinterpret_cast(cur_ptr); + cur_ptr = (char*)cur_ptr + size; + remaining_space -= size; + return result; + } else { + std::ostringstream oss; + oss << "Buffer overflow when allocating memory for " << name << " with size " << size + << " and alignment " << alignment << ", but only " << remaining_space + << " bytes available in AlignedAllocator. Increase the workspace buffer size."; + FLASHINFER_ERROR(oss.str()); + } + return nullptr; + } + + size_t aligned_alloc_offset(size_t size, size_t alignment, std::string name) { + return (char*)aligned_alloc(size, alignment, name) - (char*)base_ptr; + } + + size_t num_allocated_bytes() { return (char*)cur_ptr - (char*)base_ptr; } +}; + +} // namespace flashinfer + +#endif // FLASHINFER_ALLOCATOR_H_ diff --git a/include/flashinfer/arch_condition.h b/include/flashinfer/arch_condition.h new file mode 100644 index 0000000000000000000000000000000000000000..a7d66d154e6f5256471358b4c5625c3815d63fad --- /dev/null +++ b/include/flashinfer/arch_condition.h @@ -0,0 +1,96 @@ +/* + * Copyright (c) 2022-2025, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#ifndef FLASHINFER_ARCH_CONDITION_H_ +#define FLASHINFER_ARCH_CONDITION_H_ + +namespace flashinfer { + +namespace detail { + +#ifdef __CUDA_ARCH__ + +#ifdef __CUDA_ARCH_SPECIFIC__ +static constexpr bool isArchSpecific = true; +#else +static constexpr bool isArchSpecific = false; +#endif + +struct arch_info { + static constexpr bool mIsDevice = true; + static constexpr bool mArchSpecific = isArchSpecific; + static constexpr int mMajor = __CUDA_ARCH__ / 100; + static constexpr int mMinor = __CUDA_ARCH__ / 10 % 10; + static constexpr int mArch = __CUDA_ARCH__ / 10; +}; + +#else + +struct arch_info { + static constexpr bool mIsDevice = false; + static constexpr bool mArchSpecific = false; + static constexpr int mMajor = 0; + static constexpr int mMinor = 0; + static constexpr int mArch = 0; +}; + +#endif + +#if __CUDA_ARCH__ >= 900 +#if CUDA_VERSION >= 12090 +#if !defined(__CUDA_ARCH_SPECIFIC__) && !defined(__CUDA_ARCH_FAMILY_SPECIFIC__) +#error \ + "Compiling for SM90 or newer architectures must use Arch specific or Arch Family specific target" +#endif +#endif +#endif + +} // namespace detail + +namespace arch { + +struct is_device : std::bool_constant {}; + +struct is_arch_specific : std::bool_constant {}; + +template +struct is_match : std::bool_constant {}; + +template +struct is_major : std::bool_constant {}; + +template +struct is_compatible + : std::bool_constant::value && detail::arch_info::mArch >= Arch> {}; + +inline constexpr bool is_device_v = is_device::value; + +inline constexpr bool is_arch_specific_v = is_arch_specific::value; + +template +inline constexpr bool is_match_v = is_match::value; + +template +inline constexpr bool is_major_v = is_major::value; + +template +inline constexpr bool is_compatible_v = is_compatible::value; + +} // namespace arch + +} // namespace flashinfer + +#endif // FLASHINFER_ARCH_CONDITION_H_ diff --git a/include/flashinfer/attention/blackwell/collective/fmha_common.hpp b/include/flashinfer/attention/blackwell/collective/fmha_common.hpp new file mode 100644 index 0000000000000000000000000000000000000000..e79972512870fad1f9f0aaab5ef865bb86fd48fe --- /dev/null +++ b/include/flashinfer/attention/blackwell/collective/fmha_common.hpp @@ -0,0 +1,136 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include "cute/tensor.hpp" +#include "cutlass/arch/reg_reconfig.h" +#include "cutlass/kernel_hardware_info.h" + +namespace cutlass::fmha::collective { + +using namespace cute; + +template +CUTLASS_DEVICE auto get_local_tile_tensor(const MTensor& m_tensor, const Shape& tile_shape, + int head_idx, int offset, int seq_len) { + // (N, D, H) + auto g_offset = local_tile(m_tensor(_, _, head_idx), cute::make_shape(1, get<1>(tile_shape)), + make_coord(offset, _0{})); + auto g_sequence = + make_tensor(g_offset.data(), + make_layout(cute::make_shape(seq_len, get<1>(tile_shape)), g_offset.stride())); + auto g_tensor = local_tile(g_sequence, tile_shape, make_coord(_, _0{})); + return g_tensor; +} + +template +CUTLASS_DEVICE auto get_local_tile_t_tensor(const MTensor& m_tensor, const Shape& tile_shape, + int head_idx, int offset, int seq_len) { + // (D, N, H) + auto g_offset = local_tile(m_tensor(_, _, head_idx), cute::make_shape(get<0>(tile_shape), 1), + make_coord(_0{}, offset)); + auto g_sequence = + make_tensor(g_offset.data(), + make_layout(cute::make_shape(get<0>(tile_shape), seq_len), g_offset.stride())); + auto g_tensor = local_tile(g_offset, tile_shape, make_coord(_0{}, _)); + return g_tensor; +} + +template +CUTE_DEVICE void gemm_reset_zero_acc(Atom& atom, TA const& tA, TB const& tB, TC&& tC) { + constexpr int rA = decltype(rank(tA))::value; + constexpr int rB = decltype(rank(tB))::value; + constexpr int rC = decltype(rank(tC))::value; + static_assert(rA == 3 && rB == 3 && rC == 3); + + CUTLASS_PRAGMA_UNROLL + for (int k_block = 0; k_block < size<2>(tA); k_block++) { + cute::gemm(atom, tA(_, _, k_block), tB(_, _, k_block), tC); + atom.accumulate_ = decltype(atom.accumulate_)::One; + } +} + +template +CUTE_DEVICE void gemm_zero_acc(Atom& atom, TA const& tA, TB const& tB, TC&& tC) { + atom.accumulate_ = decltype(atom.accumulate_)::Zero; + gemm_reset_zero_acc(atom, tA, tB, tC); +} + +template +CUTE_DEVICE constexpr auto unstageSmemLayout(Layout const& layout, Stages stages = {}) { + return composition(layout, prepend(make_layout(stages), _)); +} + +template +CUTE_DEVICE T warp_uniform(T a) { + return __shfl_sync(0xffffffff, a, 0); +} + +template +CUTE_HOST_DEVICE constexpr auto to_tiled_mma_sm100_ts( + TiledMMA, + cute::C, cute::integral_constant, + cute::integral_constant, + cute::integral_constant, + cute::integral_constant>, + TAs...>, + TMs...>) { + return TiledMMA< + MMA_Atom>, + TAs...>, + TMs...>{}; +} + +template +CUTE_HOST_DEVICE constexpr auto to_tiled_mma_sm100_ts( + TiledMMA< + MMA_Atom, + TAs...>, + TMs...>) { + return TiledMMA, + TAs...>, + TMs...>{}; +} + +template +CUTLASS_DEVICE void warpgroup_reg_set() { + if constexpr (RegCount < 128) { + cutlass::arch::warpgroup_reg_dealloc(); + } else { + cutlass::arch::warpgroup_reg_alloc(); + } +} + +} // namespace cutlass::fmha::collective diff --git a/include/flashinfer/attention/blackwell/collective/fmha_fusion.hpp b/include/flashinfer/attention/blackwell/collective/fmha_fusion.hpp new file mode 100644 index 0000000000000000000000000000000000000000..1fbef6b12086200ed7d2c876f975044612fa3f14 --- /dev/null +++ b/include/flashinfer/attention/blackwell/collective/fmha_fusion.hpp @@ -0,0 +1,203 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include "cute/tensor.hpp" +#include "cutlass/cutlass.h" + +namespace cutlass::fmha::collective { + +using namespace cute; + +struct NoMask { + template + CUTLASS_DEVICE int get_trip_count(BlkCoord const& blk_coord, TileShape const& tile_shape, + ProblemSize const& problem_size) { + return ceil_div(get<1>(problem_size), get<1>(tile_shape)); + } + + template + CUTLASS_DEVICE int get_masked_trip_count(BlkCoord const& blk_coord, TileShape const& tile_shape, + ProblemSize const& problem_size) { + return 0; + } + + template + CUTLASS_DEVICE int get_unmasked_trip_count(BlkCoord const& blk_coord, TileShape const& tile_shape, + ProblemSize const& problem_size) { + return get_trip_count(blk_coord, tile_shape, problem_size); + } + + template + CUTLASS_DEVICE void apply_mask(AccQK& acc_qk, IndexQK const& index_qk, + ProblemSize const& problem_size) { + return; + } +}; + +struct ResidualMask : NoMask { + using Base = NoMask; + + template + CUTLASS_DEVICE int get_masked_trip_count(BlkCoord const& blk_coord, TileShape const& tile_shape, + ProblemSize const& problem_size) { + if (get<1>(problem_size) % get<1>(tile_shape) != 0) { + return 1; + } + return 0; + } + + template + CUTLASS_DEVICE int get_unmasked_trip_count(BlkCoord const& blk_coord, TileShape const& tile_shape, + ProblemSize const& problem_size) { + // if the sequence length does not divide the tile size evenly + if (get<1>(problem_size) % get<1>(tile_shape) != 0) { + return get_trip_count(blk_coord, tile_shape, problem_size) - 1; + } + return get_trip_count(blk_coord, tile_shape, problem_size); + } + + template + CUTLASS_DEVICE void apply_mask(AccQK& acc_qk, IndexQK const& index_qk, + ProblemSize const& problem_size) { + // This is useful is seqlen_k % kBlockN != 0 since it masks + // the remaining elements out from softmax. + // d % kHeadDim != 0 or seqlen_q % kBlockM do not suffer from similar + // issues as they are transparently taken care of by TMA and the + // epilogue, if it is instantiated with predication support. + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(acc_qk); i++) { + auto pos = index_qk(i); + if (get<1>(pos) >= get<1>(problem_size)) { + acc_qk(i) = -INFINITY; + } + } + } +}; + +struct CausalMask : NoMask { + using Base = NoMask; + + template + CUTLASS_DEVICE int get_trip_count(BlkCoord const& blk_coord, TileShape const& tile_shape, + ProblemSize const& problem_size) { + // See note below on different ways to think about causal attention + // Again, we'd add the offset_q into the max_blocks_q calculation + int offset_q = int(get<1>(problem_size)) - int(get<0>(problem_size)); + int max_blocks_k = Base::get_trip_count(blk_coord, tile_shape, problem_size); + int max_blocks_q = + ceil_div((get<0>(blk_coord) + 1) * get<0>(tile_shape) + offset_q, get<1>(tile_shape)); + return std::min(max_blocks_k, max_blocks_q); + } + + template + CUTLASS_DEVICE int get_masked_trip_count(BlkCoord const& blk_coord, TileShape const& tile_shape, + ProblemSize const& problem_size) { + return get_trip_count(blk_coord, tile_shape, problem_size); + } + + template + CUTLASS_DEVICE int get_unmasked_trip_count(BlkCoord const& blk_coord, TileShape const& tile_shape, + ProblemSize const& problem_size) { + return get_trip_count(blk_coord, tile_shape, problem_size) - + get_masked_trip_count(blk_coord, tile_shape, problem_size); + } + + template + CUTLASS_DEVICE void apply_mask(AccQK& acc_qk, IndexQK const& index_qk, + ProblemSize const& problem_size) { + // There are two ways to do causal if N_Q != N_K + // (1) is to assume that the Q is at the beginning of the matrix + // - this is what we demonstrate here + // (2) is that it is at the end of the matrix + // - this is usually what we want for inference settings + // where we only compute the next row and use cache for the rest + // - if you'd like this, you only need to add an offset like so: + // get<0>(pos) + offset_q < get<1>(pos) + int offset_q = int(get<1>(problem_size)) - int(get<0>(problem_size)); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(acc_qk); i++) { + auto pos = index_qk(i); + if ((get<0>(pos) + offset_q < get<1>(pos)) || (get<1>(pos) >= get<1>(problem_size))) { + acc_qk(i) = -INFINITY; + } + } + } +}; + +struct VariableLength { + int* segment_offsets = nullptr; + // CUTE_HOST_DEVICE operator int() const { return 0; } +}; + +template +struct is_variable_length : std::false_type {}; +template <> +struct is_variable_length : std::true_type {}; +template +constexpr bool is_variable_length_v = is_variable_length::value; + +template +CUTE_HOST_DEVICE constexpr auto apply_variable_length(Shape const& shape, Idx const& idx) { + return transform_leaf(shape, [&](auto const& s) { + if constexpr (is_variable_length_v>) { + return s.segment_offsets[idx + 1] - s.segment_offsets[idx]; + } else { + return s; + } + }); +} + +template +CUTE_HOST_DEVICE constexpr auto apply_variable_length(Shape const& shape, Coord const& coord, + Idx const& idx) { + auto new_shape = apply_variable_length(shape, idx); + auto new_coord = transform_leaf(shape, coord, [&](auto const& s, auto const& c) { + if constexpr (is_variable_length_v>) { + return cute::make_tuple(c, s.segment_offsets[idx]); + } else { + return c; + } + }); + return cute::make_tuple(new_shape, new_coord); +} + +} // namespace cutlass::fmha::collective + +namespace cute { + +template <> +struct is_integral : true_type {}; + +CUTE_HOST_DEVICE +void print(cutlass::fmha::collective::VariableLength a) { printf("Varlen<%p>", a.segment_offsets); } + +} // namespace cute diff --git a/include/flashinfer/attention/blackwell/collective/sm100_fmha_fwd_epilogue_tma_warpspecialized.hpp b/include/flashinfer/attention/blackwell/collective/sm100_fmha_fwd_epilogue_tma_warpspecialized.hpp new file mode 100644 index 0000000000000000000000000000000000000000..d6ab631f2ccb7b5cd8dd30ae19b63702479a39e7 --- /dev/null +++ b/include/flashinfer/attention/blackwell/collective/sm100_fmha_fwd_epilogue_tma_warpspecialized.hpp @@ -0,0 +1,180 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include "cute/layout.hpp" +#include "cutlass/cutlass.h" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "fmha_common.hpp" + +namespace cutlass::fmha::collective { + +template +struct Sm100FmhaFwdEpilogueTmaWarpspecialized { + using ElementOut = Element; + using Pipeline = cutlass::PipelineAsync<2>; + // using ShapeT = cute::Shape>; + // using StrideO = cute::Shape>; + // using LayoutO = cute::Layout; + using ShapeT = cute::Shape, int32_t>>; + using StrideO = cute::Shape, int32_t>>; + using LayoutO = cute::Layout; + + using ShapeLSE = cute::Shape>; + using StrideLSE = cute::Shape>; + using LayoutLSE = cute::Layout; + + // using SmemLayoutO = decltypa(make_layout(append<3>(select<0,1>(TileShape_WG{}), _2{}))); + using SmemLayoutAtomO = decltype(cutlass::gemm::collective::detail::sm100_smem_selector< + cute::UMMA::Major::K, Element, tuple_element_t<0, TileShape>, + tuple_element_t<1, TileShape>>()); + // using SmemLayoutAtomO = decltype(make_ordered_layout(select<0,1>(TileShape{}), Step<_1, + // _0>{})); + using SmemLayoutO = + decltype(tile_to_shape(SmemLayoutAtomO{}, replace<2>(TileShape{}, _2{}), Step<_2, _1, _3>{})); + using SmemLayoutO_ = SmemLayoutO; + + struct TensorStorage { + using SmemLayoutO = SmemLayoutO_; + cute::array_aligned> smem_o; + }; + struct Arguments { + Element* ptr_O; + LayoutO layout_O; + + ElementAcc* ptr_LSE; + LayoutLSE layout_LSE; + int max_qo_len; + }; + + using TMA_O = decltype(make_tma_copy( + SM90_TMA_STORE{}, make_tensor((Element*)nullptr, repeat_like(StrideO{}, 0), StrideO{}), + SmemLayoutO{}(_, _, _0{}))); + + struct Params { + TMA_O tma_store_o; + LayoutO layout_O; + ElementAcc* ptr_LSE; + LayoutLSE layout_LSE; + int max_qo_len; + }; + + template + static Params to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, + void* workspace = nullptr) { + static_assert(is_variable_length_v>); + auto ptr_O = args.ptr_O; + LayoutO layout_O = args.layout_O; + + auto tma_store_o = + make_tma_copy(SM90_TMA_STORE{}, make_tensor(ptr_O, layout_O), SmemLayoutO{}(_, _, _0{})); + + return {tma_store_o, layout_O, args.ptr_LSE, args.layout_LSE, args.max_qo_len}; + } + + CUTLASS_DEVICE + static void prefetch_tma_descriptors(Params const& params) { + cute::prefetch_tma_descriptor(params.tma_store_o.get_tma_descriptor()); + } + + const Params& params; + + CUTLASS_DEVICE Sm100FmhaFwdEpilogueTmaWarpspecialized(const Params& params) : params(params) {} + + template + CUTLASS_DEVICE auto store(BlkCoord const& blk_coord, ProblemShape const& problem_shape, + Params const& params, ParamsProblemShape const& params_problem_shape, + TensorStorage& shared_storage, Pipeline& pipeline, + typename Pipeline::PipelineState& pipeline_consumer_state) { + int qo_tile_idx = get<0>(blk_coord); + int qo_head_idx = get<2, 0>(blk_coord); + int batch_idx = get<2, 1>(blk_coord); + int qo_len = get<0>(problem_shape); + int qo_segment_offset = get<0>(params_problem_shape).segment_offsets[batch_idx]; + uint32_t lane_predicate = cute::elect_one_sync(); + + using X = Underscore; + + int o0_index = 2 * get<0>(blk_coord); + int o1_index = 2 * get<0>(blk_coord) + 1; + + int offs_0 = params.max_qo_len - qo_len; + int offs_2_1 = qo_segment_offset + qo_len; + BlkCoord blk_coord_updated = blk_coord; + get<2, 1>(blk_coord_updated) = 0; + + Tensor mO = params.tma_store_o.get_tma_tensor(params.layout_O.shape()); + + Tensor mO_qdl = domain_offset(make_coord(offs_0, _0{}, make_coord(_0{}, offs_2_1)), mO); + + Tensor gO_qdl = local_tile(mO_qdl, TileShape{}, make_coord(_, _, _), Step<_1, _1, X>{}); + Tensor gO = gO_qdl(_, _, _, _0{}, get<2>(blk_coord_updated)); + + Tensor sO = make_tensor(make_smem_ptr(shared_storage.smem_o.data()), SmemLayoutO{}); + auto block_tma = params.tma_store_o.get_slice(0); + Tensor tOsO = block_tma.partition_S(sO); + Tensor tOgO = block_tma.partition_D(gO); + + auto pipeline_release_state = pipeline_consumer_state; + + // O1 O2 + // one pipeline: O + // wait from corr, issue tma store on smem + pipeline.consumer_wait(pipeline_consumer_state); + ++pipeline_consumer_state; + + if (lane_predicate) { + copy(params.tma_store_o, tOsO(_, _, _, _0{}), tOgO(_, _, _, o0_index)); + } + tma_store_arrive(); + + pipeline.consumer_wait(pipeline_consumer_state); + ++pipeline_consumer_state; + + if (lane_predicate) { + copy(params.tma_store_o, tOsO(_, _, _, _1{}), tOgO(_, _, _, o1_index)); + } + tma_store_arrive(); + + tma_store_wait<1>(); + + pipeline.consumer_release(pipeline_release_state); + ++pipeline_release_state; + + tma_store_wait<0>(); + + pipeline.consumer_release(pipeline_release_state); + ++pipeline_release_state; + } +}; + +} // namespace cutlass::fmha::collective diff --git a/include/flashinfer/attention/blackwell/collective/sm100_fmha_fwd_mainloop_tma_warpspecialized.hpp b/include/flashinfer/attention/blackwell/collective/sm100_fmha_fwd_mainloop_tma_warpspecialized.hpp new file mode 100644 index 0000000000000000000000000000000000000000..9393f66079d8f8f1860464fbc89b4b1f19946a9e --- /dev/null +++ b/include/flashinfer/attention/blackwell/collective/sm100_fmha_fwd_mainloop_tma_warpspecialized.hpp @@ -0,0 +1,1159 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include "cute/arch/simd_sm100.hpp" +#include "cute/layout.hpp" +#include "cute/tensor.hpp" +#include "cutlass/arch/memory_sm80.h" +#include "cutlass/cutlass.h" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "fmha_common.hpp" +#include "fmha_fusion.hpp" +#include "sm100_fmha_load_tma_warpspecialized.hpp" + +namespace cutlass::fmha::collective { + +using namespace cute; + +template > +struct Sm100FmhaFwdMainloopTmaWarpspecialized { + using Element = Element_; + using ElementQK = ElementQK_; + using ElementPV = ElementPV_; + using TileShape = decltype(select<0, 1>(TileShapeQK_{})); + using TileShapeQK = decltype(shape_div(TileShapeQK_{}, ThreadShape{})); + using TileShapePV = decltype(shape_div(TileShapePV_{}, ThreadShape{})); + using StrideQ = StrideQ_; + using StrideK = StrideK_; + using StrideV = StrideV_; + using Mask = Mask_; + + static constexpr int StageCountQ = 2; + static constexpr int StageCountKV = + get<2>(TileShapeQK{}) == 128 ? 2 : 1; // sizeof(Element_) == 1 ? 2 : 2; + + using StagesQ = cutlass::gemm::collective::StageCount; + using StagesKV = cutlass::gemm::collective::StageCount; + + using ClusterShape = Shape<_1, _1, _1>; + + static const int Alignment = 128 / sizeof_bits_v; + + using CollectiveMmaQK = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, Element, StrideQ, Alignment, Element, + StrideK, Alignment, ElementQK, TileShapeQK, ClusterShape, + cutlass::gemm::collective::StageCount<3> /* we change it later anyways*/, + cutlass::gemm::KernelTmaWarpSpecialized1SmSm100>::CollectiveOp; + + using CollectiveMmaPV = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + // the stride for A does not matter since we do not load from smem at all + Element, StrideK, Alignment, Element, StrideV, Alignment, ElementPV, TileShapePV, + ClusterShape, cutlass::gemm::collective::StageCount<3> /* we change it later anyways*/, + cutlass::gemm::KernelTmaWarpSpecialized1SmSm100>::CollectiveOp; + + using SmemLayoutQ = + decltype(unstageSmemLayout(typename CollectiveMmaQK::SmemLayoutA{}, Int{})); + using SmemLayoutK = + decltype(unstageSmemLayout(typename CollectiveMmaQK::SmemLayoutB{}, Int{})); + using SmemLayoutV = + decltype(unstageSmemLayout(typename CollectiveMmaPV::SmemLayoutB{}, Int{})); + + struct TensorStorage { + cute::array_aligned> smem_q; + cute::array_aligned> smem_k; + cute::array_aligned> smem_v; + }; + + enum class TmemAllocation : uint32_t { + kSizeS = 128, + kSizeO = 128, + kSizeP = 32, + S0 = 0, + S1 = S0 + kSizeS, // 128 + V0 = S0, // 0 // stats storage from softmax to correction + V1 = S1, // 128 + P0 = S0 + kSizeP, // 32 + P1 = S1 + kSizeP, // 160 + O0 = S1 + kSizeS, // 256 + O1 = O0 + kSizeO, // 384 + kEnd = O1 + kSizeO // 512 + }; + + // indices for V0 / V1 + enum : int { kIdxOldRowMax = 0, kIdxNewRowMax = 1, kIdxFinalRowSum = 0, kIdxFinalRowMax = 1 }; + + // from load to mma warp, protects q in smem + using PipelineQ = + cutlass::PipelineTmaUmmaAsync; + + // from load to mma warp, protects k/v in smem + using PipelineK = + cutlass::PipelineTmaUmmaAsync; + + using PipelineV = + cutlass::PipelineTmaUmmaAsync; + + // from mma to softmax0/1 warp, protects S in tmem + // (not sure yet about the reverse direction) + // there is one pipe per softmax warp, and the mma warp alternates between them + using PipelineS = cutlass::PipelineUmmaAsync<1>; + + // from softmax0/1/ to correction wg + using PipelineC = cutlass::PipelineAsync<1>; + + // from mma to correction + using PipelineO = cutlass::PipelineUmmaAsync<2>; + + // from corr to epilogue + using PipelineE = cutlass::PipelineAsync<2>; + + using OrderBarrierSoftmax = cutlass::OrderedSequenceBarrier< + /*stages*/ 1, /*groups*/ 2>; + + static const int TransactionBytesLoadQ = + cutlass::bits_to_bytes(cosize(take<0, 3>(SmemLayoutQ{})) * cute::sizeof_bits_v); + + static const int TransactionBytesLoadK = + cutlass::bits_to_bytes(cosize(take<0, 3>(SmemLayoutK{})) * cute::sizeof_bits_v); + + static const int TransactionBytesLoadV = + cutlass::bits_to_bytes(cosize(take<0, 3>(SmemLayoutV{})) * cute::sizeof_bits_v); + + // static_assert( + // cutlass::bits_to_bytes(cosize(take<0, 3>(SmemLayoutK{})) * cute::sizeof_bits_v) == + // cutlass::bits_to_bytes(cosize(take<0, 3>(SmemLayoutV{})) * + // cute::sizeof_bits_v), + // "K and V smem layouts must be of equal size"); + + using Load = Sm100FmhaLoadTmaWarpspecialized; + using LayoutQ = typename Load::LayoutQ; + using LayoutK = typename Load::LayoutK; + using LayoutV = typename Load::LayoutV; + + struct Arguments { + typename Load::Arguments load; + + float scale_softmax; + + // scaling factors to dequantize QKV + float scale_q = 1.0f; + float scale_k = 1.0f; + float scale_v = 1.0f; + + // scaling factor to quantize O + float inv_scale_o = 1.0f; + }; + + struct Params { + typename Load::Params load; + + float scale_softmax; + float scale_softmax_log2; + + float scale_output; + }; + + template + static bool can_implement(ProblemShape const& problem_shape, Arguments const& args) { + return true; + } + + template + static Params to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, + void* workspace) { + float scale_softmax = args.scale_softmax; + float log2_e = static_cast(std::log2(std::exp(1.0))); + + return Params{Load::to_underlying_arguments(problem_shape, args.load, workspace), + args.scale_q * args.scale_k * scale_softmax, + args.scale_q * args.scale_k * log2_e * scale_softmax, + args.scale_v * args.inv_scale_o}; + } + + CUTLASS_DEVICE + static void prefetch_tma_descriptors(Params const& params) { + Load::prefetch_tma_descriptors(params.load); + } + + template + CUTLASS_DEVICE void load(BlkCoord const& blk_coord, ProblemShape const& problem_shape, + Params const& params, ParamsProblemShape const& params_problem_shape, + TensorStorage& storage, PipelineQ& pipeline_q, + typename PipelineQ::PipelineState& pipeline_q_producer_state, + PipelineK& pipeline_k, + typename PipelineK::PipelineState& pipeline_k_producer_state, + PipelineV& pipeline_v, + typename PipelineV::PipelineState& pipeline_v_producer_state) { + Load load; + load.load(blk_coord, problem_shape, params.load, params_problem_shape, storage, pipeline_q, + pipeline_q_producer_state, pipeline_k, pipeline_k_producer_state, pipeline_v, + pipeline_v_producer_state); + } + + template + CUTLASS_DEVICE auto mma( + BlkCoord const& blk_coord, Params const& params, ProblemShape const& problem_shape, + TensorStorage& storage, PipelineQ& pipeline_q, + typename PipelineQ::PipelineState& pipeline_q_consumer_state, PipelineK& pipeline_k, + typename PipelineK::PipelineState& pipeline_k_consumer_state, PipelineV& pipeline_v, + typename PipelineV::PipelineState& pipeline_v_consumer_state, PipelineS& pipeline_s0, + typename PipelineS::PipelineState& pipeline_s0_producer_state, PipelineS& pipeline_s1, + typename PipelineS::PipelineState& pipeline_s1_producer_state, PipelineO& pipeline_corr, + typename PipelineO::PipelineState& pipeline_corr_producer_state) { + auto pipeline_q_release_state = pipeline_q_consumer_state; + auto pipeline_k_release_state = pipeline_k_consumer_state; + auto pipeline_v_release_state = pipeline_v_consumer_state; + + int mask_tile_count = Mask{}.get_trip_count(blk_coord, TileShape{}, problem_shape); + + typename CollectiveMmaQK::TiledMma mma_qk; + ThrMMA thr_mma_qk = mma_qk.get_slice(0); + + typename CollectiveMmaPV::TiledMma mma_pv; + TiledMMA mma_pv_ts = to_tiled_mma_sm100_ts(mma_pv); + ThrMMA thr_mma_pv = mma_pv_ts.get_slice(0); + + Tensor sQ = make_tensor(make_smem_ptr(storage.smem_q.data()), SmemLayoutQ{}); + Tensor sK = make_tensor(make_smem_ptr(storage.smem_k.data()), SmemLayoutK{}); + Tensor sV = make_tensor(make_smem_ptr(storage.smem_v.data()), SmemLayoutV{}); + + Tensor tSrQ = thr_mma_qk.make_fragment_A(sQ); + Tensor tSrK = thr_mma_qk.make_fragment_B(sK); + Tensor tOrV = thr_mma_pv.make_fragment_B(sV); + + // tmem layout is + // S0 S1`O0 O1 + // sequential in memory, where S overlaps with P and V + + Tensor tStS = partition_fragment_C(mma_qk, select<0, 1>(TileShapeQK{})); + Tensor tOtO = partition_fragment_C(mma_pv_ts, select<0, 1>(TileShapePV{})); + + Tensor tStS0 = tStS; + tStS0.data() = tStS.data().get() + uint32_t(TmemAllocation::S0); + Tensor tStS1 = tStS; + tStS1.data() = tStS.data().get() + uint32_t(TmemAllocation::S1); + + Tensor tOtO0 = tOtO; + tOtO0.data() = tOtO.data().get() + uint32_t(TmemAllocation::O0); + Tensor tOtO1 = tOtO; + tOtO1.data() = tOtO.data().get() + uint32_t(TmemAllocation::O1); + + Tensor sP = + make_tensor(make_smem_ptr((Element*)nullptr), typename CollectiveMmaPV::SmemLayoutA{}); + Tensor tOrP = thr_mma_pv.make_fragment_A(sP)(_, _, _, _0{}); // slice out staging + + Tensor tOrP0 = tOrP; + tOrP0.data() = tOrP0.data().get() + uint32_t(TmemAllocation::P0); + Tensor tOrP1 = tOrP; + tOrP1.data() = tOrP1.data().get() + uint32_t(TmemAllocation::P1); + + int k_index = 0; + int v_index = 0; + int q_index = 0; + + // wait for Q1 + q_index = pipeline_q_consumer_state.index(); + pipeline_q.consumer_wait(pipeline_q_consumer_state); + ++pipeline_q_consumer_state; + + Tensor tSrQ0 = tSrQ(_, _, _, q_index); + + // wait for K1 + k_index = pipeline_k_consumer_state.index(); + pipeline_k.consumer_wait(pipeline_k_consumer_state); + ++pipeline_k_consumer_state; + + // gemm Q1 * K1 -> S1 + pipeline_s0.producer_acquire(pipeline_s0_producer_state); + + gemm_zero_acc(mma_qk, tSrQ0, tSrK(_, _, _, k_index), tStS0); + + pipeline_s0.producer_commit(pipeline_s0_producer_state); + ++pipeline_s0_producer_state; + + // release K1 + if constexpr (get<1>(ThreadShape{}) > 1) { + pipeline_k.consumer_release(pipeline_k_release_state); + ++pipeline_k_release_state; + } + + // wait for Q2 + if constexpr (get<0>(ThreadShape{}) > 1 || get<2>(ThreadShape{}) > 1) { + q_index = pipeline_q_consumer_state.index(); + pipeline_q.consumer_wait(pipeline_q_consumer_state); + ++pipeline_q_consumer_state; + } + + Tensor tSrQ1 = tSrQ(_, _, _, q_index); + + if constexpr (get<1>(ThreadShape{}) > 1) { + k_index = pipeline_k_consumer_state.index(); + pipeline_k.consumer_wait(pipeline_k_consumer_state); + ++pipeline_k_consumer_state; + } + + pipeline_s1.producer_acquire(pipeline_s1_producer_state); + + // gemm Q2 * K1 -> S2 + gemm_zero_acc(mma_qk, tSrQ1, tSrK(_, _, _, k_index), tStS1); + + pipeline_s1.producer_commit(pipeline_s1_producer_state); + ++pipeline_s1_producer_state; + + // release K1 + pipeline_k.consumer_release(pipeline_k_release_state); + ++pipeline_k_release_state; + + // wait for V1 + v_index = pipeline_v_consumer_state.index(); + pipeline_v.consumer_wait(pipeline_v_consumer_state); + ++pipeline_v_consumer_state; + + // this acquire returns the ownership of all of S0 to the mma warp + // including the P0 part + // acquire corr first to take it out of the critical + // path since softmax takes longer + pipeline_corr.producer_acquire(pipeline_corr_producer_state); + pipeline_s0.producer_acquire(pipeline_s0_producer_state); + + // gemm P1 * V1 -> O1 + gemm_zero_acc(mma_pv_ts, tOrP0, tOrV(_, _, _, v_index), tOtO0); + + pipeline_corr.producer_commit(pipeline_corr_producer_state); + ++pipeline_corr_producer_state; + + if constexpr (get<1>(ThreadShape{}) > 1) { + pipeline_v.consumer_release(pipeline_v_release_state); + ++pipeline_v_release_state; + } + + mma_pv_ts.accumulate_ = UMMA::ScaleOut::Zero; + + // loop: + mask_tile_count -= 1; + for (; mask_tile_count > 0; mask_tile_count -= 1) { + // wait for Ki + k_index = (pipeline_k_consumer_state.index()); + pipeline_k.consumer_wait(pipeline_k_consumer_state); + ++pipeline_k_consumer_state; + + // gemm Q1 * Ki -> S1 + gemm_zero_acc(mma_qk, tSrQ0, tSrK(_, _, _, k_index), tStS0); + + pipeline_s0.producer_commit(pipeline_s0_producer_state); + ++pipeline_s0_producer_state; + + if constexpr (get<1>(ThreadShape{}) > 1) { + pipeline_k.consumer_release(pipeline_k_release_state); + ++pipeline_k_release_state; + } + + // gemm P2 * V(i-1) -> O2 + if constexpr (get<1>(ThreadShape{}) > 1) { + v_index = pipeline_v_consumer_state.index(); + pipeline_v.consumer_wait(pipeline_v_consumer_state); + ++pipeline_v_consumer_state; + } + + pipeline_corr.producer_acquire(pipeline_corr_producer_state); + pipeline_s1.producer_acquire(pipeline_s1_producer_state); + + gemm_reset_zero_acc(mma_pv_ts, tOrP1, tOrV(_, _, _, v_index), tOtO1); + + pipeline_corr.producer_commit(pipeline_corr_producer_state); + ++pipeline_corr_producer_state; + + // release V(i-1) + pipeline_v.consumer_release(pipeline_v_release_state); + ++pipeline_v_release_state; + + if constexpr (get<1>(ThreadShape{}) > 1) { + k_index = (pipeline_k_consumer_state.index()); + pipeline_k.consumer_wait(pipeline_k_consumer_state); + ++pipeline_k_consumer_state; + } + + // gemm Q2 * Ki -> S2 + gemm_zero_acc(mma_qk, tSrQ1, tSrK(_, _, _, k_index), tStS1); + + pipeline_s1.producer_commit(pipeline_s1_producer_state); + ++pipeline_s1_producer_state; + + // release Ki + pipeline_k.consumer_release(pipeline_k_release_state); + ++pipeline_k_release_state; + + // wait for Vi + v_index = (pipeline_v_consumer_state.index()); + pipeline_v.consumer_wait(pipeline_v_consumer_state); + ++pipeline_v_consumer_state; + + // gemm P1 * Vi -> O1 + pipeline_corr.producer_acquire(pipeline_corr_producer_state); + + pipeline_s0.producer_acquire(pipeline_s0_producer_state); + + gemm_reset_zero_acc(mma_pv_ts, tOrP0, tOrV(_, _, _, v_index), tOtO0); + + pipeline_corr.producer_commit(pipeline_corr_producer_state); + ++pipeline_corr_producer_state; + + if constexpr (get<1>(ThreadShape{}) > 1) { + pipeline_v.consumer_release(pipeline_v_release_state); + ++pipeline_v_release_state; + } + } + + // release Q1 + pipeline_q.consumer_release(pipeline_q_release_state); + ++pipeline_q_release_state; + + // release Q2 + if constexpr (get<0>(ThreadShape{}) > 1) { + pipeline_q.consumer_release(pipeline_q_release_state); + ++pipeline_q_release_state; + } + + // wait for Vi + if constexpr (get<1>(ThreadShape{}) > 1) { + v_index = pipeline_v_consumer_state.index(); + pipeline_v.consumer_wait(pipeline_v_consumer_state); + ++pipeline_v_consumer_state; + } + + // gemm P2 * Vi -> O2 + pipeline_corr.producer_acquire(pipeline_corr_producer_state); + pipeline_s1.producer_acquire(pipeline_s1_producer_state); + + gemm_reset_zero_acc(mma_pv_ts, tOrP1, tOrV(_, _, _, v_index), tOtO1); + + pipeline_corr.producer_commit(pipeline_corr_producer_state); + ++pipeline_corr_producer_state; + + // release Vi + pipeline_v.consumer_release(pipeline_v_release_state); + ++pipeline_v_release_state; + + pipeline_s0.producer_commit(pipeline_s0_producer_state); + ++pipeline_s0_producer_state; + + pipeline_s1.producer_commit(pipeline_s1_producer_state); + ++pipeline_s1_producer_state; + + // T0 S00 B1, T0 S10 B1, T0 S00 B2, T0 S01 B1, T0 S10 B2, T0 S11 B1, T0 S01 B2, T1 S00 B1, T0 + // S11 B2, ... Q1 * K1 , Q2 * K1 , S11 * V1 , Q1 * K2 , S21 * V1 , Q2 * K2 , S12 * V2 , Q1 * + // K3 , S22 * K2 , ... + } + + template + CUTLASS_DEVICE auto softmax_step(float& row_max, float& row_sum, Stage stage, bool final_call, + BlkCoord const& blk_coord, CountingTensor const& cS, + Params const& params, ProblemShape const& problem_shape, + PipelineS& pipeline_s, + typename PipelineS::PipelineState& pipeline_s_consumer_state, + PipelineC& pipeline_c, + typename PipelineC::PipelineState& pipeline_c_producer_state, + OrderBarrierSoftmax& order_s) { + Tensor tScS = typename CollectiveMmaQK::TiledMma{}.get_slice(0).partition_C(cS); + + Tensor tStS = + partition_fragment_C(typename CollectiveMmaQK::TiledMma{}, select<0, 1>(TileShapeQK{})); + tStS.data() = uint32_t(stage == _0{} ? TmemAllocation::S0 : TmemAllocation::S1); + + Tensor tStS_v = tStS.compose(make_layout(make_shape(_128{}, _2{}))); + tStS_v.data() = uint32_t(stage == _0{} ? TmemAllocation::V0 : TmemAllocation::V1); + Tensor tScS_v = tScS.compose(make_layout(make_shape(_128{}, _2{}))); + + auto tilePlikeFP32 = get<1>(TileShapeQK{}) / Int{} * Int{}; + Tensor tStS_P = tStS.compose(make_layout(make_shape(_128{}, tilePlikeFP32))); + tStS_P.data() = warp_uniform(uint32_t(stage == _0{} ? TmemAllocation::P0 : TmemAllocation::P1)); + Tensor tScS_P = tScS.compose(make_layout(make_shape(_128{}, tilePlikeFP32))); + + // Each thread owns a single row + using TMEM_LOAD = SM100_TMEM_LOAD_32dp32b32x; // 4x32 threads with 128 cols of 32b elem + using TMEM_STORE = SM100_TMEM_STORE_32dp32b32x; // 4x32 threads with 128 cols of 8b elem + using TMEM_STORE_V = SM100_TMEM_STORE_32dp32b2x; // 4x32 threads with 2 cols of 32b elem + + int thread_idx = threadIdx.x % (4 * cutlass::NumThreadsPerWarp); + + auto tiled_tmem_load = make_tmem_copy(TMEM_LOAD{}, tStS); + auto thr_tmem_load = tiled_tmem_load.get_slice(thread_idx); + + Tensor tTMEM_LOADtS = thr_tmem_load.partition_S(tStS); + Tensor tTMEM_LOADcS = thr_tmem_load.partition_D(tScS); + + auto tiled_tmem_storev = make_tmem_copy(TMEM_STORE_V{}, tStS_v); + auto thr_tmem_storev = tiled_tmem_storev.get_slice(thread_idx); + + Tensor tTMEM_STOREVtS = thr_tmem_storev.partition_D(tStS_v); + Tensor tTMEM_STOREVcS = thr_tmem_storev.partition_S(tScS_v); + + auto tiled_tmem_store = make_tmem_copy(TMEM_STORE{}, tStS_P); + auto thr_tmem_store = tiled_tmem_store.get_slice(thread_idx); + + Tensor tTMEM_STOREtS_x4 = thr_tmem_store.partition_D(tStS_P); + tTMEM_STOREtS_x4.data() = warp_uniform(tTMEM_STOREtS_x4.data().get()); + Tensor tTMEM_STOREcS = thr_tmem_store.partition_S(tScS_P); + + // wait on tensor core pipe + pipeline_s.consumer_wait(pipeline_s_consumer_state); + + // read all of S from tmem into reg mem + Tensor tTMEM_LOADrS = make_tensor(shape(tTMEM_LOADcS)); + copy(tiled_tmem_load, tTMEM_LOADtS, tTMEM_LOADrS); + + if constexpr (need_apply_mask) { + Mask{}.apply_mask(tTMEM_LOADrS, tTMEM_LOADcS, problem_shape); + } + + ElementQK old_row_max = row_max; + { + // compute rowmax + float row_max_0 = row_max; + float row_max_1 = row_max; + float row_max_2 = row_max; + float row_max_3 = row_max; + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(tTMEM_LOADrS); i += 4) { + row_max_0 = ::fmax(row_max_0, tTMEM_LOADrS(i)); + row_max_1 = ::fmax(row_max_1, tTMEM_LOADrS(i + 1)); + row_max_2 = ::fmax(row_max_2, tTMEM_LOADrS(i + 2)); + row_max_3 = ::fmax(row_max_3, tTMEM_LOADrS(i + 3)); + } + row_max = ::fmax(row_max_0, row_max_1); + row_max = ::fmax(row_max, row_max_2); + row_max = ::fmax(row_max, row_max_3); + } + + ElementQK row_max_safe = row_max == -INFINITY ? 0 : row_max; + + Tensor tTMEM_STOREVrS = make_tensor(shape(tTMEM_STOREVcS)); + tTMEM_STOREVrS(kIdxOldRowMax) = old_row_max; + tTMEM_STOREVrS(kIdxNewRowMax) = row_max_safe; + copy(tiled_tmem_storev, tTMEM_STOREVrS, tTMEM_STOREVtS); + + pipeline_c.producer_commit(pipeline_c_producer_state); + ++pipeline_c_producer_state; + + // notify correction wg that they are ready (might need addtl ordering between S0 and S1 WG's) + + ElementQK scale = params.scale_softmax_log2; + ElementQK row_max_scale = row_max_safe * scale; + + float2 scale_fp32x2 = make_float2(scale, scale); + float2 minus_row_max_scale_fp32x2 = make_float2(-row_max_scale, -row_max_scale); + + Tensor tTMEM_STORErS_x4 = make_tensor(shape(tTMEM_STOREcS)); + + constexpr int kConversionsPerStep = 2; + + Tensor tTMEM_STORErS_x4_e = recast>(tTMEM_STORErS_x4); + + NumericArrayConverter convert; + + const int kReleasePipeCount = 10; // must be multiple of 2 + + order_s.wait(); + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(tTMEM_LOADrS); i += 2) { + float2 in = make_float2(tTMEM_LOADrS(i + 0), tTMEM_LOADrS(i + 1)); + float2 out; + cute::fma(out, scale_fp32x2, in, minus_row_max_scale_fp32x2); + tTMEM_LOADrS(i + 0) = out.x; + tTMEM_LOADrS(i + 1) = out.y; + + tTMEM_LOADrS(i + 0) = ::exp2f(tTMEM_LOADrS(i + 0)); + tTMEM_LOADrS(i + 1) = ::exp2f(tTMEM_LOADrS(i + 1)); + + Array in_conv; + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < kConversionsPerStep; j++) { + in_conv[j] = tTMEM_LOADrS(i + j); + } + tTMEM_STORErS_x4_e[i / kConversionsPerStep] = convert(in_conv); + + if (i == size(tTMEM_LOADrS) - kReleasePipeCount) { + order_s.arrive(); + } + + // this prevents register spills in fp16 + if constexpr (size<2>(tTMEM_STORErS_x4) == _2{}) { + if (i == size(tTMEM_LOADrS) - 6) { + copy(tiled_tmem_store, tTMEM_STORErS_x4(_, _, 0), tTMEM_STOREtS_x4(_, _, 0)); + } + } + } + + // tmem_store(reg_S8) -> op_P + CUTE_STATIC_ASSERT_V(size<2>(tTMEM_STORErS_x4) <= _2{}); + CUTE_STATIC_ASSERT_V(size<1>(tTMEM_STORErS_x4) == _1{}); + copy(tiled_tmem_store, tTMEM_STORErS_x4(_, _, size<2>(tTMEM_STORErS_x4) - 1), + tTMEM_STOREtS_x4(_, _, size<2>(tTMEM_STORErS_x4) - 1)); + + cutlass::arch::fence_view_async_tmem_store(); + + // notify tensor core warp that P is ready + pipeline_s.consumer_release(pipeline_s_consumer_state); + ++pipeline_s_consumer_state; + + pipeline_c.producer_acquire(pipeline_c_producer_state); + + ElementQK acc_scale = 0.5f * ::exp2f(scale * (old_row_max - row_max_safe)); + row_sum *= acc_scale; + // row_sum = sum(reg_S) + float2 local_row_sum_f32x2 = make_float2(row_sum, row_sum); + float2 local_row_sum_1 = make_float2(0, 0); + float2 local_row_sum_2 = make_float2(0, 0); + float2 local_row_sum_3 = make_float2(0, 0); + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(tTMEM_LOADrS); i += 8) { + // row_sum += tTMEM_LOADrS(i); + float2 in = make_float2(tTMEM_LOADrS(i), tTMEM_LOADrS(i + 1)); + cute::add(local_row_sum_f32x2, local_row_sum_f32x2, in); + + in = make_float2(tTMEM_LOADrS(i + 2), tTMEM_LOADrS(i + 2 + 1)); + cute::add(local_row_sum_1, local_row_sum_1, in); + + in = make_float2(tTMEM_LOADrS(i + 4), tTMEM_LOADrS(i + 4 + 1)); + cute::add(local_row_sum_2, local_row_sum_2, in); + + in = make_float2(tTMEM_LOADrS(i + 6), tTMEM_LOADrS(i + 6 + 1)); + cute::add(local_row_sum_3, local_row_sum_3, in); + } + + cute::add(local_row_sum_f32x2, local_row_sum_f32x2, local_row_sum_1); + cute::add(local_row_sum_2, local_row_sum_2, local_row_sum_3); + cute::add(local_row_sum_f32x2, local_row_sum_f32x2, local_row_sum_2); + float local_row_sum = local_row_sum_f32x2.x + local_row_sum_f32x2.y; + + row_sum = local_row_sum; + + if (final_call) { + // re-acquire the S part in the final step + pipeline_s.consumer_wait(pipeline_s_consumer_state); + + Tensor tTMEM_STOREVrS = make_tensor(shape(tTMEM_STOREVcS)); + tTMEM_STOREVrS(kIdxFinalRowMax) = row_max; + tTMEM_STOREVrS(kIdxFinalRowSum) = row_sum; + copy(tiled_tmem_storev, tTMEM_STOREVrS, tTMEM_STOREVtS); + } + } + + template + CUTLASS_DEVICE auto softmax(Stage stage, BlkCoord const& blk_coord, Params const& params, + ProblemShape const& problem_shape, PipelineS& pipeline_s, + typename PipelineS::PipelineState& pipeline_s_consumer_state, + PipelineC& pipeline_c, + typename PipelineC::PipelineState& pipeline_c_producer_state, + OrderBarrierSoftmax& order_s) { + int mask_tile_count = Mask{}.get_unmasked_trip_count(blk_coord, TileShape{}, problem_shape); + + ElementQK row_max = -INFINITY; + ElementQK row_sum = 0; + + Tensor cS_base = make_identity_tensor(select<0, 1>(TileShapeQK{})); + auto logical_offset = make_coord(get<0>(blk_coord) * get<0>(TileShape{}) + + (stage % get<0>(ThreadShape{})) * get<0>(TileShapeQK{}), + 0 + (stage % get<1>(ThreadShape{})) * get<1>(TileShapeQK{})); + Tensor cS = domain_offset(logical_offset, cS_base); + + pipeline_c.producer_acquire(pipeline_c_producer_state); + + CUTLASS_PRAGMA_NO_UNROLL + for (; mask_tile_count > 0; mask_tile_count -= 1) { + softmax_step( + row_max, row_sum, stage, + (mask_tile_count == 1) && + (Mask{}.get_masked_trip_count(blk_coord, TileShape{}, problem_shape) == 0), + blk_coord, cS, params, problem_shape, pipeline_s, pipeline_s_consumer_state, pipeline_c, + pipeline_c_producer_state, order_s); + + cS.data() = cS.data() + E<1>{} * get<1>(ThreadShape{}) * get<1>(TileShapeQK{}); + } + + // Masked iterations + mask_tile_count = Mask{}.get_masked_trip_count(blk_coord, TileShape{}, problem_shape); + + CUTLASS_PRAGMA_NO_UNROLL + for (; mask_tile_count > 0; mask_tile_count -= 1) { + softmax_step( + row_max, row_sum, stage, mask_tile_count == 1, blk_coord, cS, params, problem_shape, + pipeline_s, pipeline_s_consumer_state, pipeline_c, pipeline_c_producer_state, order_s); + + cS.data() = cS.data() + E<1>{} * get<1>(ThreadShape{}) * get<1>(TileShapeQK{}); + } + + pipeline_c.producer_commit(pipeline_c_producer_state); + ++pipeline_c_producer_state; + + pipeline_c.producer_acquire(pipeline_c_producer_state); + // empty step to sync against pipe s + pipeline_s.consumer_release(pipeline_s_consumer_state); + ++pipeline_s_consumer_state; + } + + template + CUTLASS_DEVICE auto correction_epilogue(float scale, Stage stage, TensorO const& sO_01) { + using ElementOut = typename TensorO::value_type; + + int thread_idx = threadIdx.x % (4 * cutlass::NumThreadsPerWarp); + + Tensor sO = sO_01(_, _, stage); + + // As opposed to the softmax, we do not have enough registers here + // to load all of the values (for tile kv = 128), so we loop + // good values would be either 32 or 64 + const int kCorrectionTileSize = 32 / sizeof(ElementOut); + + using TMEM_LOAD = + std::conditional_t; // 4x32 threads with 64 cols of 32b elem + + typename CollectiveMmaPV::TiledMma mma; + Tensor cO = make_identity_tensor(select<0, 1>(TileShapePV{})); + Tensor tOtO = partition_fragment_C(mma, select<0, 1>(TileShapePV{})); + Tensor tOcO = mma.get_slice(0).partition_C(cO); + Tensor tOsO = mma.get_slice(0).partition_C(sO); + + Tensor tOtO_i = + logical_divide(tOtO, make_layout(make_shape(_128{}, Int{}))); + Tensor tOcO_i = + logical_divide(tOcO, make_layout(make_shape(_128{}, Int{}))); + Tensor tOsO_i = + logical_divide(tOsO, make_layout(make_shape(_128{}, Int{}))); + + if constexpr (decltype(stage == _0{})::value) { + tOtO_i.data() = tOtO_i.data().get() + uint32_t(TmemAllocation::O0); + } else { + static_assert(decltype(stage == _1{})::value, "stage is either 0 or 1"); + tOtO_i.data() = tOtO_i.data().get() + uint32_t(TmemAllocation::O1); + } + + auto tiled_tmem_load = make_tmem_copy(TMEM_LOAD{}, tOtO_i(make_coord(_, _), _0{})); + auto thr_tmem_load = tiled_tmem_load.get_slice(thread_idx); + + Tensor tTMEM_LOADtO = thr_tmem_load.partition_S(tOtO_i(make_coord(_, _), _)); + Tensor tTMEM_LOADcO = thr_tmem_load.partition_D(tOcO_i(make_coord(_, _), _)); + Tensor tTMEM_LOADsO = thr_tmem_load.partition_D(tOsO_i(make_coord(_, _), _)); + + float2 scale_f32x2 = make_float2(scale, scale); + + // loop: + // TMEM_LOAD, FMUL2 scale, TMEM_STORE + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < get<1>(TileShapePV{}) / kCorrectionTileSize; i++) { + Tensor tTMEM_LOADtO_i = tTMEM_LOADtO(_, _0{}, _0{}, i); + Tensor tTMEM_LOADsO_i = tTMEM_LOADsO(_, _0{}, _0{}, i); + + Tensor tTMrO = make_tensor(shape(tTMEM_LOADcO(_, _0{}, _0{}, i))); + + copy(tiled_tmem_load, tTMEM_LOADtO_i, tTMrO); + +#ifndef ONLY_SOFTMAX + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < size(tTMrO); j += 2) { + float2 in = make_float2(tTMrO(j), tTMrO(j + 1)); + float2 out; + cute::mul(out, scale_f32x2, in); + tTMrO(j) = out.x; + tTMrO(j + 1) = out.y; + } +#endif + + constexpr int N = 4 / sizeof(ElementOut); + NumericArrayConverter convert; + + Tensor tSMrO = make_tensor_like(tTMrO); + + Tensor tCs = recast(tTMrO); + Tensor tCd = recast(tSMrO); + + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < size(tCs); j++) { + tCd(j) = convert.convert(tCs(j)); + } + + Tensor tSMsO_i = recast(tTMEM_LOADsO_i); + Tensor tSMrO_i = recast(tSMrO); + + copy(AutoVectorizingCopyWithAssumedAlignment<128>{}, tSMrO_i, tSMsO_i); + } + + cutlass::arch::fence_view_async_shared(); + } + + CUTLASS_DEVICE auto correction_rescale(float scale, uint32_t tmem_O) { + int thread_idx = threadIdx.x % (4 * cutlass::NumThreadsPerWarp); + + // As opposed to the softmax, we do not have enough registers here + // to load all of the values (for tile kv = 128), so we loop + // good values would be either 32 or 64 + const int kCorrectionTileSize = 16; + + using TMEM_LOAD = SM100_TMEM_LOAD_32dp32b16x; // 4x32 threads with 64 cols of 32b elem + using TMEM_STORE = SM100_TMEM_STORE_32dp32b16x; // 4x32 threads with 64 cols of 32b elem + + typename CollectiveMmaPV::TiledMma mma; + Tensor cO = make_identity_tensor(select<0, 1>(TileShapePV{})); + Tensor tOtO = partition_fragment_C(mma, select<0, 1>(TileShapePV{})); + Tensor tOcO = mma.get_slice(0).partition_C(cO); + + Tensor tOtO_i = tOtO.compose(make_layout(make_shape(_128{}, Int{}))); + Tensor tOcO_i = tOcO.compose(make_layout(make_shape(_128{}, Int{}))); + + tOtO_i.data() = tOtO_i.data().get() + tmem_O; + + auto tiled_tmem_load = make_tmem_copy(TMEM_LOAD{}, tOtO_i); + auto thr_tmem_load = tiled_tmem_load.get_slice(thread_idx); + auto tiled_tmem_store = make_tmem_copy(TMEM_STORE{}, tOtO_i); + auto thr_tmem_store = tiled_tmem_store.get_slice(thread_idx); + + Tensor tTMEM_LOADtO = thr_tmem_load.partition_S(tOtO_i); + Tensor tTMEM_LOADcO = thr_tmem_load.partition_D(tOcO_i); + Tensor tTMEM_STOREtO = thr_tmem_store.partition_D(tOtO_i); + Tensor tTMEM_STOREcO = thr_tmem_store.partition_S(tOcO_i); + static_assert(shape(tTMEM_STOREcO) == shape(tTMEM_LOADcO)); + + float2 scale_f32x2 = make_float2(scale, scale); + + Tensor tTMrO = + make_tensor(make_shape(shape(tTMEM_LOADcO), Int<128 / kCorrectionTileSize>{})); + + auto copy_in = [&](int i) { + Tensor tTMEM_LOADtO_i = tTMEM_LOADtO; + tTMEM_LOADtO_i.data() = tTMEM_LOADtO_i.data().get() + uint32_t(i * kCorrectionTileSize); + Tensor tTMrO_i = tTMrO(_, i).compose(make_layout(shape<0>(tTMrO))); + copy(tiled_tmem_load, tTMEM_LOADtO_i, tTMrO_i); + }; + + auto copy_out = [&](int i) { + Tensor tTMEM_STOREtO_i = tTMEM_STOREtO; + tTMEM_STOREtO_i.data() = tTMEM_STOREtO_i.data().get() + uint32_t(i * kCorrectionTileSize); + Tensor tTMrO_i = tTMrO(_, i).compose(make_layout(shape<0>(tTMrO))); + copy(tiled_tmem_store, tTMrO_i, tTMEM_STOREtO_i); + }; + + // sequence: LLMSLMSLMSS + + // loop: + // TMEM_LOAD, FMUL2 scale, TMEM_STORE + copy_in(0); + + int count = get<1>(TileShapePV{}) / kCorrectionTileSize; + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < count; i++) { + if (i != count - 1) { + copy_in(i + 1); + } + + Tensor tTMrO_i = tTMrO(_, i).compose(make_layout(shape<0>(tTMrO))); + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < size(tTMrO_i); j += 2) { + float2 in = make_float2(tTMrO_i(j), tTMrO_i(j + 1)); + float2 out; + cute::mul(out, scale_f32x2, in); + tTMrO_i(j) = out.x; + tTMrO_i(j + 1) = out.y; + } + + copy_out(i); + } + } + + template + CUTLASS_DEVICE auto correction( + BlkCoord const& blk_coord, Params const& params, + ParamsProblemShape const& params_problem_shape, ProblemShape const& problem_shape, + TensorStorageEpi& shared_storage_epi, PipelineC& pipeline_s0_c, + typename PipelineC::PipelineState& pipeline_s0_c_consumer_state, PipelineC& pipeline_s1_c, + typename PipelineC::PipelineState& pipeline_s1_c_consumer_state, PipelineO& pipeline_o, + typename PipelineO::PipelineState& pipeline_o_consumer_state, PipelineE& pipeline_epi, + typename PipelineE::PipelineState& pipeline_epi_producer_state, + CollectiveEpilogue& epilogue) { + int mask_tile_count = Mask{}.get_trip_count(blk_coord, TileShape{}, problem_shape); + + int thread_idx = threadIdx.x % (4 * cutlass::NumThreadsPerWarp); + + Tensor tStS = + partition_fragment_C(typename CollectiveMmaQK::TiledMma{}, select<0, 1>(TileShapeQK{})); + + Tensor cS = make_identity_tensor(select<0, 1>(TileShapeQK{})); + Tensor tScS = typename CollectiveMmaQK::TiledMma{}.get_slice(0).partition_C(cS); + + Tensor tStS_v = tStS.compose(make_layout(make_shape(_128{}, _2{}))); + Tensor tScS_v = tScS.compose(make_layout(make_shape(_128{}, _2{}))); + + using TMEM_LOAD_V = SM100_TMEM_LOAD_32dp32b2x; // 4x32 threads with 2 cols of 32b elem + + auto tiled_tmem_loadv = make_tmem_copy(TMEM_LOAD_V{}, tStS_v); + auto thr_tmem_loadv = tiled_tmem_loadv.get_slice(thread_idx); + + Tensor tTMEM_LOADVtS = thr_tmem_loadv.partition_S(tStS_v); + Tensor tTMEM_LOADVcS = thr_tmem_loadv.partition_D(tScS_v); + + Tensor tTMEM_LOADVtS0 = tTMEM_LOADVtS; + tTMEM_LOADVtS0.data() = tTMEM_LOADVtS0.data().get() + uint32_t(TmemAllocation::V0); + Tensor tTMEM_LOADVtS1 = tTMEM_LOADVtS; + tTMEM_LOADVtS1.data() = tTMEM_LOADVtS1.data().get() + uint32_t(TmemAllocation::V1); + + // ignore first signal from softmax as no correction is required + pipeline_s0_c.consumer_wait(pipeline_s0_c_consumer_state); + pipeline_s0_c.consumer_release(pipeline_s0_c_consumer_state); + ++pipeline_s0_c_consumer_state; + + pipeline_s1_c.consumer_wait(pipeline_s1_c_consumer_state); + + // handle the last iteration differently (i.e. tmem_load/stsm for epi) + mask_tile_count -= 1; + + CUTLASS_PRAGMA_NO_UNROLL + for (; mask_tile_count > 0; mask_tile_count -= 1) { + pipeline_s0_c.consumer_wait(pipeline_s0_c_consumer_state); + + Tensor tTMEM_LOADVrS = make_tensor(shape(tTMEM_LOADVcS)); + + // read row_wise new global max + copy(tiled_tmem_loadv, tTMEM_LOADVtS0, tTMEM_LOADVrS); + + // e^(scale * (old_max - new_max) + float scale = ::exp2f(params.scale_softmax_log2 * + (tTMEM_LOADVrS(kIdxOldRowMax) - tTMEM_LOADVrS(kIdxNewRowMax))); + + pipeline_o.consumer_wait(pipeline_o_consumer_state); + + correction_rescale(scale, uint32_t(TmemAllocation::O0)); + + pipeline_s1_c.consumer_release(pipeline_s1_c_consumer_state); + ++pipeline_s1_c_consumer_state; + + cutlass::arch::fence_view_async_tmem_store(); + + pipeline_o.consumer_release(pipeline_o_consumer_state); + ++pipeline_o_consumer_state; + + pipeline_s1_c.consumer_wait(pipeline_s1_c_consumer_state); + + copy(tiled_tmem_loadv, tTMEM_LOADVtS1, tTMEM_LOADVrS); + + scale = ::exp2f(params.scale_softmax_log2 * + (tTMEM_LOADVrS(kIdxOldRowMax) - tTMEM_LOADVrS(kIdxNewRowMax))); + + pipeline_o.consumer_wait(pipeline_o_consumer_state); + + correction_rescale(scale, uint32_t(TmemAllocation::O1)); + + pipeline_s0_c.consumer_release(pipeline_s0_c_consumer_state); + ++pipeline_s0_c_consumer_state; + + cutlass::arch::fence_view_async_tmem_store(); + + pipeline_o.consumer_release(pipeline_o_consumer_state); + ++pipeline_o_consumer_state; + } + + pipeline_s1_c.consumer_release(pipeline_s1_c_consumer_state); + ++pipeline_s1_c_consumer_state; + + // do the final correction to O1 + // better to somehow special-case it in the loop above + // doesn't matter for non-persistent code, but if it were + // persistent we do not want to release O too early + + pipeline_s0_c.consumer_wait(pipeline_s0_c_consumer_state); + + // read from V0 + // read row_sum and final row_max here + Tensor tTMEM_LOADVrS = make_tensor(shape(tTMEM_LOADVcS)); + copy(tiled_tmem_loadv, tTMEM_LOADVtS0, tTMEM_LOADVrS); + + pipeline_s0_c.consumer_release(pipeline_s0_c_consumer_state); + ++pipeline_s0_c_consumer_state; + + pipeline_o.consumer_wait(pipeline_o_consumer_state); + pipeline_epi.producer_acquire(pipeline_epi_producer_state); + // store to epi smem + + // loop: + // TMEM_LOAD + // FMUL2 scale = 1 / global_sum * out_quant_scale + // F2FP + // store to smem + Tensor sO = make_tensor(make_smem_ptr(shared_storage_epi.smem_o.data()), + typename TensorStorageEpi::SmemLayoutO{}); + Tensor gLSE = make_tensor(make_gmem_ptr(epilogue.params.ptr_LSE), epilogue.params.layout_LSE); + correction_epilogue(params.scale_output / tTMEM_LOADVrS(kIdxFinalRowSum), _0{}, sO); + if (epilogue.params.ptr_LSE != nullptr) { + int qo_tile_idx = get<0>(blk_coord); + int qo_head_idx = get<2, 0>(blk_coord); + int batch_idx = get<2, 1>(blk_coord); + int qo_len = get<0>(problem_shape); + int segment_offset = get<0>(params_problem_shape).segment_offsets[batch_idx]; + int row_idx = get<0>(tTMEM_LOADVcS(_0{})) + get<0>(TileShape{}) * qo_tile_idx; + + ElementPV lse = __log2f(tTMEM_LOADVrS(kIdxFinalRowSum)) + + params.scale_softmax_log2 * tTMEM_LOADVrS(kIdxFinalRowMax); + + if (row_idx < qo_len) { + gLSE(segment_offset + row_idx, qo_head_idx) = lse; + } + } + // correction_epilogue(params.scale_output, _0{}, sO); + + cutlass::arch::fence_view_async_tmem_load(); + + pipeline_o.consumer_release(pipeline_o_consumer_state); + ++pipeline_o_consumer_state; + + pipeline_epi.producer_commit(pipeline_epi_producer_state); + ++pipeline_epi_producer_state; + + pipeline_s1_c.consumer_wait(pipeline_s1_c_consumer_state); + + // load from V1 + copy(tiled_tmem_loadv, tTMEM_LOADVtS1, tTMEM_LOADVrS); + + pipeline_s1_c.consumer_release(pipeline_s1_c_consumer_state); + ++pipeline_s1_c_consumer_state; + + pipeline_o.consumer_wait(pipeline_o_consumer_state); + pipeline_epi.producer_acquire(pipeline_epi_producer_state); + + correction_epilogue(params.scale_output / tTMEM_LOADVrS(kIdxFinalRowSum), _1{}, sO); + + if (epilogue.params.ptr_LSE != nullptr) { + int qo_tile_idx = get<0>(blk_coord); + int qo_head_idx = get<2, 0>(blk_coord); + int batch_idx = get<2, 1>(blk_coord); + int qo_len = get<0>(problem_shape); + int segment_offset = get<0>(params_problem_shape).segment_offsets[batch_idx]; + int row_idx = + get<0>(tTMEM_LOADVcS(_0{})) + get<0>(TileShape{}) * qo_tile_idx + get<0>(TileShapeQK{}); + + ElementPV lse = __log2f(tTMEM_LOADVrS(kIdxFinalRowSum)) + + params.scale_softmax_log2 * tTMEM_LOADVrS(kIdxFinalRowMax); + + if (row_idx < qo_len) { + gLSE(segment_offset + row_idx, qo_head_idx) = lse; + } + } + // correction_epilogue(params.scale_output, _1{}, sO); + cutlass::arch::fence_view_async_tmem_load(); + + pipeline_o.consumer_release(pipeline_o_consumer_state); + ++pipeline_o_consumer_state; + + pipeline_epi.producer_commit(pipeline_epi_producer_state); + ++pipeline_epi_producer_state; + } + + template + CUTLASS_DEVICE auto correction_empty( + BlkCoord const& blk_coord, Params const& params, ProblemShape const& problem_shape, + ParamsProblemShape const& params_problem_shape, TensorStorageEpi& shared_storage_epi, + PipelineE& pipeline_epi, typename PipelineE::PipelineState& pipeline_epi_producer_state, + CollectiveEpilogue& epilogue) { + pipeline_epi.producer_acquire(pipeline_epi_producer_state); + + Tensor sO = make_tensor(make_smem_ptr(shared_storage_epi.smem_o.data()), + typename TensorStorageEpi::SmemLayoutO{}); + Tensor gLSE = make_tensor(make_gmem_ptr(epilogue.params.ptr_LSE), epilogue.params.layout_LSE); + int thread_idx = threadIdx.x % (4 * NumThreadsPerWarp); + + using ElementOut = typename CollectiveEpilogue::ElementOut; + auto tiled_copy = make_cotiled_copy( + Copy_Atom, ElementOut>{}, + make_ordered_layout(make_shape(_128{}, Int{}), + Step<_1, _0>{}), + sO.layout()); + + auto thr_copy = tiled_copy.get_slice(thread_idx); + auto tOgO = thr_copy.partition_D(sO); + auto tOrO = make_tensor(shape(tOgO(_, _, _, _0{}))); + clear(tOrO); + + copy(tiled_copy, tOrO, tOgO(_, _, _, _0{})); + + if (epilogue.params.ptr_LSE != nullptr) { + int qo_tile_idx = get<0>(blk_coord); + int qo_head_idx = get<2, 0>(blk_coord); + int batch_idx = get<2, 1>(blk_coord); + int qo_len = get<0>(problem_shape); + int segment_offset = get<0>(params_problem_shape).segment_offsets[batch_idx]; + int row_idx = thread_idx + get<0>(TileShape{}) * qo_tile_idx; + + if (row_idx < qo_len) { + gLSE(segment_offset + row_idx, qo_head_idx) = -cuda::std::numeric_limits::infinity(); + } + } + + pipeline_epi.producer_commit(pipeline_epi_producer_state); + ++pipeline_epi_producer_state; + + copy(tiled_copy, tOrO, tOgO(_, _, _, _1{})); + cutlass::arch::fence_view_async_shared(); + pipeline_epi.producer_acquire(pipeline_epi_producer_state); + + if (epilogue.params.ptr_LSE != nullptr) { + int qo_tile_idx = get<0>(blk_coord); + int qo_head_idx = get<2, 0>(blk_coord); + int batch_idx = get<2, 1>(blk_coord); + int qo_len = get<0>(problem_shape); + int segment_offset = get<0>(params_problem_shape).segment_offsets[batch_idx]; + int row_idx = thread_idx + get<0>(TileShape{}) * qo_tile_idx + get<0>(TileShapeQK{}); + + if (row_idx < qo_len) { + gLSE(segment_offset + row_idx, qo_head_idx) = -cuda::std::numeric_limits::infinity(); + } + } + + cutlass::arch::fence_view_async_shared(); + pipeline_epi.producer_commit(pipeline_epi_producer_state); + ++pipeline_epi_producer_state; + } +}; + +} // namespace cutlass::fmha::collective diff --git a/include/flashinfer/attention/blackwell/collective/sm100_fmha_gen_epilogue_warpspecialized.hpp b/include/flashinfer/attention/blackwell/collective/sm100_fmha_gen_epilogue_warpspecialized.hpp new file mode 100644 index 0000000000000000000000000000000000000000..4eb97b540b81c8a36669a9fbdc956e87dde1496c --- /dev/null +++ b/include/flashinfer/attention/blackwell/collective/sm100_fmha_gen_epilogue_warpspecialized.hpp @@ -0,0 +1,82 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include "cute/layout.hpp" +#include "cutlass/cutlass.h" + +namespace cutlass::fmha::collective { + +template +struct Sm100FmhaGenEpilogueWarpspecialized { + using Pipeline = cutlass::PipelineAsync<2>; + + using SmemLayoutO = Layout>; + using SmemLayoutO_ = SmemLayoutO; + using Element = Element_; + using StrideOOrig = StrideO_; + using StrideO = decltype(replace<0>(StrideOOrig{}, 0)); + + struct TensorStorage { + using SmemLayoutO = SmemLayoutO_; + cute::array_aligned> smem_o; + }; + + struct Arguments { + Element* ptr_o; + StrideO dO; + }; + + using Params = Arguments; + + const Params& params; + + CUTLASS_DEVICE Sm100FmhaGenEpilogueWarpspecialized(const Params& params) : params(params) {} + + template + static Params to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, + void* workspace = nullptr) { + return args; + } + + CUTLASS_DEVICE + static void prefetch_tma_descriptors(Params const& params) { /* no-op */ } + + template + CUTLASS_DEVICE auto store(BlkCoord const& blk_coord_in, ProblemShape const& problem_shape, + Params const& params, ParamsProblemShape const& params_problem_shape, + TensorStorage& shared_storage, Pipeline& pipeline, + typename Pipeline::PipelineState& pipeline_consumer_state) { + /* no-op */ + } +}; + +} // namespace cutlass::fmha::collective diff --git a/include/flashinfer/attention/blackwell/collective/sm100_fmha_gen_mainloop_warpspecialized.hpp b/include/flashinfer/attention/blackwell/collective/sm100_fmha_gen_mainloop_warpspecialized.hpp new file mode 100644 index 0000000000000000000000000000000000000000..8a96c1cd0f88e3316c84a43f75c3acf8605931cb --- /dev/null +++ b/include/flashinfer/attention/blackwell/collective/sm100_fmha_gen_mainloop_warpspecialized.hpp @@ -0,0 +1,1064 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include "collective/fmha_common.hpp" +#include "collective/fmha_fusion.hpp" +#include "collective/sm100_fmha_load_cpasync_warpspecialized.hpp" +#include "cute/arch/simd_sm100.hpp" +#include "cute/layout.hpp" +#include "cute/tensor.hpp" +#include "cutlass/arch/memory_sm80.h" +#include "cutlass/cutlass.h" +#include "cutlass/gemm/collective/collective_builder.hpp" + +namespace cutlass::fmha::collective { + +using namespace cute; + +template > +struct Sm100FmhaGenMainloopWarpspecialized { + using Element = Element_; + using ElementQK = ElementQK_; + using ElementPV = ElementPV_; + using ElementAcc = ElementPV_; + using ElementOut = ElementOut_; + using TileShape = TileShape_; + using StrideQOrig = StrideQ_; + using StrideQ = decltype(replace<0>(StrideQ_{}, 0)); + using StrideNewK = StrideNewK_; + using StrideNewV = StrideNewV_; + using StrideCacheK = StrideK_; + using StrideCacheV = StrideV_; + using StrideK = StrideK_; + using StrideV = StrideV_; + using StrideOOrig = StrideO_; + using StrideO = decltype(replace<0>(StrideO_{}, 0)); + using Mask = Mask_; + + static constexpr int StageCountQ = get<1>(TileShape{}) == 256 ? 1 : 2; + static constexpr int StageCountKV = 256 * 11 / get<1>(TileShape{}); + + using StagesQ = cutlass::gemm::collective::StageCount; + using StagesKV = cutlass::gemm::collective::StageCount; + + using ClusterShape = Shape<_1, _1, _1>; + + static const int Alignment = 128 / sizeof_bits_v; + + using TileShapeQK = decltype(shape_div(TileShape{}, ThreadShape{})); + + using TileShapePV = decltype(select<0, 2, 1>(TileShapeQK{})); + + using CollectiveMmaQK = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, Element, StrideQ, Alignment, Element, + StrideK, Alignment, ElementQK, TileShapeQK, ClusterShape, + cutlass::gemm::collective::StageCount<3> /* we change it later anyways*/, + cutlass::gemm::KernelTmaWarpSpecialized1SmSm100>::CollectiveOp; + + using CollectiveMmaPV = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, + // the stride for A does not matter since we do not load from smem at all + Element, StrideK, Alignment, Element, decltype(select<1, 0, 2>(StrideV{})), Alignment, + ElementPV, TileShapePV, ClusterShape, + cutlass::gemm::collective::StageCount<3> /* we change it later anyways*/, + cutlass::gemm::KernelTmaWarpSpecialized1SmSm100>::CollectiveOp; + + using SmemLayoutQ = + decltype(unstageSmemLayout(typename CollectiveMmaQK::SmemLayoutA{}, Int{})); + using SmemLayoutK = + decltype(unstageSmemLayout(typename CollectiveMmaQK::SmemLayoutB{}, Int{})); + using SmemLayoutV = + decltype(unstageSmemLayout(typename CollectiveMmaPV::SmemLayoutB{}, Int{})); + + struct TensorStorage { + cute::array_aligned> smem_q; + union { + cute::array_aligned> smem_k; + cute::array_aligned> smem_v; + }; + }; + + enum class TmemAllocation : uint32_t { + kSizeS = 128, + kSizeO = 128, + kSizeP = 32, + S0 = 0, + S1 = S0 + kSizeS, + V0 = S0, // stats storage from softmax to correction + V1 = S1, + P0 = S0 + kSizeP, + P1 = S1 + kSizeP, + O0 = S1 + kSizeS, + O1 = O0 + kSizeO, + kEnd = O1 + kSizeO + }; + + // indices for V0 / V1 + enum : int { kIdxOldRowMax = 0, kIdxNewRowMax = 1, kIdxFinalRowSum = 0, kIdxFinalRowMax = 1 }; + + // from load to mma warp, protects q in smem + using PipelineQ = + cutlass::PipelineUmmaConsumerAsync; + + // from load to mma warp, protects k/v in smem + using PipelineKV = + cutlass::PipelineUmmaConsumerAsync; + + // from mma to softmax0/1 warp, protects S in tmem + // (not sure yet about the reverse direction) + // there is one pipe per softmax warp, and the mma warp alternates between them + using PipelineS = cutlass::PipelineUmmaAsync<1>; + + // from softmax0/1/ to correction wg + using PipelineC = cutlass::PipelineAsync<1>; + + // from mma to correction + using PipelineO = cutlass::PipelineUmmaAsync<2>; + + // from corr to epilogue + using PipelineE = cutlass::PipelineAsync<2>; + + using OrderBarrierSoftmax = cutlass::OrderedSequenceBarrier< + /*stages*/ 1, /*groups*/ 2>; + + static_assert( + cutlass::bits_to_bytes(cosize(take<0, 3>(SmemLayoutK{})) * cute::sizeof_bits_v) == + cutlass::bits_to_bytes(cosize(take<0, 3>(SmemLayoutV{})) * cute::sizeof_bits_v), + "K and V smem layouts must be of equal size"); + + using Load = + Sm100FmhaLoadCpAsyncWarpspecialized; + + struct Arguments { + typename Load::Arguments load; + + // if zero, defaults to 1/sqrt(D) + float scale_softmax = 0.0f; + + // scaling factors to dequantize QKV + float scale_q = 1.0f; + float scale_k = 1.0f; + float scale_v = 1.0f; + + // scaling factor to quantize O + float inv_scale_o = 1.0f; + }; + + struct Params { + typename Load::Params load; + + float scale_softmax; + float scale_softmax_log2; + + float scale_output; + }; + + template + static bool can_implement(ProblemShape const& problem_shape, Arguments const& args) { + return true; + } + + template + static Params to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, + void* workspace) { + float scale_softmax = args.scale_softmax; + if (scale_softmax == 0.0f) { + scale_softmax = 1.0f / (float)std::sqrt(get<2>(problem_shape)); + } + float log2_e = static_cast(std::log2(std::exp(1.0))); + + return Params{Load::to_underlying_arguments(problem_shape, args.load, workspace), + args.scale_q * args.scale_k * scale_softmax, + args.scale_q * args.scale_k * log2_e * scale_softmax, + args.scale_v * args.inv_scale_o}; + } + + CUTLASS_DEVICE + static void prefetch_tma_descriptors(Params const& params) { + Load::prefetch_tma_descriptors(params.load); + } + + template + CUTLASS_DEVICE void load(BlkCoord const& blk_coord, ProblemShape const& problem_shape, + Params const& params, ParamsProblemShape const& params_problem_shape, + TensorStorage& storage, PipelineQ& pipeline_q, + typename PipelineQ::PipelineState& pipeline_q_producer_state, + PipelineKV& pipeline_kv, + typename PipelineKV::PipelineState& pipeline_kv_producer_state) { + Load load; + load.load(blk_coord, problem_shape, params.load, params_problem_shape, storage, pipeline_q, + pipeline_q_producer_state, pipeline_kv, pipeline_kv_producer_state); + } + + template + CUTLASS_DEVICE auto mma( + BlkCoord const& blk_coord, Params const& params, ProblemShape const& problem_shape, + TensorStorage& storage, PipelineQ& pipeline_q, + typename PipelineQ::PipelineState& pipeline_q_consumer_state, PipelineKV& pipeline_kv, + typename PipelineKV::PipelineState& pipeline_kv_consumer_state, PipelineS& pipeline_s0, + typename PipelineS::PipelineState& pipeline_s0_producer_state, PipelineS& pipeline_s1, + typename PipelineS::PipelineState& pipeline_s1_producer_state, PipelineO& pipeline_corr, + typename PipelineO::PipelineState& pipeline_corr_producer_state) { + auto pipeline_q_release_state = pipeline_q_consumer_state; + auto pipeline_kv_release_state = pipeline_kv_consumer_state; + + int mask_tile_count = Mask{}.get_trip_count(blk_coord, TileShape{}, problem_shape); + + typename CollectiveMmaQK::TiledMma mma_qk; + ThrMMA thr_mma_qk = mma_qk.get_slice(0); + + typename CollectiveMmaPV::TiledMma mma_pv; + TiledMMA mma_pv_ts = to_tiled_mma_sm100_ts(mma_pv); + ThrMMA thr_mma_pv = mma_pv_ts.get_slice(0); + + Tensor sQ = make_tensor(make_smem_ptr(storage.smem_q.data()), SmemLayoutQ{}); + Tensor sK = make_tensor(make_smem_ptr(storage.smem_k.data()), SmemLayoutK{}); + Tensor sV = make_tensor(make_smem_ptr(storage.smem_v.data()), SmemLayoutV{}); + + Tensor tSrQ = thr_mma_qk.make_fragment_A(sQ); + Tensor tSrK = thr_mma_qk.make_fragment_B(sK); + Tensor tOrV = thr_mma_pv.make_fragment_B(sV); + + // tmem layout is + // S0 S1`O0 O1 + // sequential in memory, where S overlaps with P and V + + Tensor tStS = partition_fragment_C(mma_qk, select<0, 1>(TileShapeQK{})); + Tensor tOtO = partition_fragment_C(mma_pv_ts, select<0, 1>(TileShapePV{})); + + Tensor tStS0 = tStS; + tStS0.data() = tStS.data().get() + uint32_t(TmemAllocation::S0); + Tensor tStS1 = tStS; + tStS1.data() = tStS.data().get() + uint32_t(TmemAllocation::S1); + + Tensor tOtO0 = tOtO; + tOtO0.data() = tOtO.data().get() + uint32_t(TmemAllocation::O0); + Tensor tOtO1 = tOtO; + tOtO1.data() = tOtO.data().get() + uint32_t(TmemAllocation::O1); + + Tensor sP = + make_tensor(make_smem_ptr((Element*)nullptr), typename CollectiveMmaPV::SmemLayoutA{}); + Tensor tOrP = thr_mma_pv.make_fragment_A(sP)(_, _, _, _0{}); // slice out staging + + Tensor tOrP0 = tOrP; + tOrP0.data() = tOrP0.data().get() + uint32_t(TmemAllocation::P0); + Tensor tOrP1 = tOrP; + tOrP1.data() = tOrP1.data().get() + uint32_t(TmemAllocation::P1); + + int k_index = 0; + int v_index = 0; + int q_index = 0; + + // wait for Q1 + q_index = pipeline_q_consumer_state.index(); + pipeline_q.consumer_wait(pipeline_q_consumer_state); + ++pipeline_q_consumer_state; + + Tensor tSrQ0 = tSrQ(_, _, _, q_index); + + // wait for K1 + k_index = pipeline_kv_consumer_state.index(); + pipeline_kv.consumer_wait(pipeline_kv_consumer_state); + ++pipeline_kv_consumer_state; + + // gemm Q1 * K1 -> S1 + pipeline_s0.producer_acquire(pipeline_s0_producer_state); + + gemm_zero_acc(mma_qk, tSrQ0, tSrK(_, _, _, k_index), tStS0); + + pipeline_s0.producer_commit(pipeline_s0_producer_state); + ++pipeline_s0_producer_state; + + // release K1 + if constexpr (get<1>(ThreadShape{}) > 1) { + pipeline_kv.consumer_release(pipeline_kv_release_state); + ++pipeline_kv_release_state; + } + + // wait for Q2 + if constexpr (get<0>(ThreadShape{}) > 1 || get<2>(ThreadShape{}) > 1) { + q_index = pipeline_q_consumer_state.index(); + pipeline_q.consumer_wait(pipeline_q_consumer_state); + ++pipeline_q_consumer_state; + } + + Tensor tSrQ1 = tSrQ(_, _, _, q_index); + + if constexpr (get<1>(ThreadShape{}) > 1) { + k_index = pipeline_kv_consumer_state.index(); + pipeline_kv.consumer_wait(pipeline_kv_consumer_state); + ++pipeline_kv_consumer_state; + } + + pipeline_s1.producer_acquire(pipeline_s1_producer_state); + + // gemm Q2 * K1 -> S2 + gemm_zero_acc(mma_qk, tSrQ1, tSrK(_, _, _, k_index), tStS1); + + pipeline_s1.producer_commit(pipeline_s1_producer_state); + ++pipeline_s1_producer_state; + + // release K1 + pipeline_kv.consumer_release(pipeline_kv_release_state); + ++pipeline_kv_release_state; + + // wait for V1 + v_index = pipeline_kv_consumer_state.index(); + pipeline_kv.consumer_wait(pipeline_kv_consumer_state); + ++pipeline_kv_consumer_state; + + // this acquire returns the ownership of all of S0 to the mma warp + // including the P0 part + // acquire corr first to take it out of the critical + // path since softmax takes longer + pipeline_corr.producer_acquire(pipeline_corr_producer_state); + pipeline_s0.producer_acquire(pipeline_s0_producer_state); + + // gemm P1 * V1 -> O1 + gemm_zero_acc(mma_pv_ts, tOrP0, tOrV(_, _, _, v_index), tOtO0); + + pipeline_corr.producer_commit(pipeline_corr_producer_state); + ++pipeline_corr_producer_state; + + if constexpr (get<1>(ThreadShape{}) > 1) { + pipeline_kv.consumer_release(pipeline_kv_release_state); + ++pipeline_kv_release_state; + } + + mma_pv_ts.accumulate_ = UMMA::ScaleOut::Zero; + + // loop: + mask_tile_count -= 1; + for (; mask_tile_count > 0; mask_tile_count -= 1) { + // wait for Ki + k_index = (pipeline_kv_consumer_state.index()); + pipeline_kv.consumer_wait(pipeline_kv_consumer_state); + ++pipeline_kv_consumer_state; + + // gemm Q1 * Ki -> S1 + gemm_zero_acc(mma_qk, tSrQ0, tSrK(_, _, _, k_index), tStS0); + + pipeline_s0.producer_commit(pipeline_s0_producer_state); + ++pipeline_s0_producer_state; + + if constexpr (get<1>(ThreadShape{}) > 1) { + pipeline_kv.consumer_release(pipeline_kv_release_state); + ++pipeline_kv_release_state; + } + + // gemm P2 * V(i-1) -> O2 + if constexpr (get<1>(ThreadShape{}) > 1) { + v_index = pipeline_kv_consumer_state.index(); + pipeline_kv.consumer_wait(pipeline_kv_consumer_state); + ++pipeline_kv_consumer_state; + } + + pipeline_corr.producer_acquire(pipeline_corr_producer_state); + pipeline_s1.producer_acquire(pipeline_s1_producer_state); + + gemm_reset_zero_acc(mma_pv_ts, tOrP1, tOrV(_, _, _, v_index), tOtO1); + + pipeline_corr.producer_commit(pipeline_corr_producer_state); + ++pipeline_corr_producer_state; + + // release V(i-1) + pipeline_kv.consumer_release(pipeline_kv_release_state); + ++pipeline_kv_release_state; + + if constexpr (get<1>(ThreadShape{}) > 1) { + k_index = (pipeline_kv_consumer_state.index()); + pipeline_kv.consumer_wait(pipeline_kv_consumer_state); + ++pipeline_kv_consumer_state; + } + + // gemm Q2 * Ki -> S2 + gemm_zero_acc(mma_qk, tSrQ1, tSrK(_, _, _, k_index), tStS1); + + pipeline_s1.producer_commit(pipeline_s1_producer_state); + ++pipeline_s1_producer_state; + + // release Ki + pipeline_kv.consumer_release(pipeline_kv_release_state); + ++pipeline_kv_release_state; + + // wait for Vi + v_index = (pipeline_kv_consumer_state.index()); + pipeline_kv.consumer_wait(pipeline_kv_consumer_state); + ++pipeline_kv_consumer_state; + + // gemm P1 * Vi -> O1 + pipeline_corr.producer_acquire(pipeline_corr_producer_state); + + pipeline_s0.producer_acquire(pipeline_s0_producer_state); + + gemm_reset_zero_acc(mma_pv_ts, tOrP0, tOrV(_, _, _, v_index), tOtO0); + + pipeline_corr.producer_commit(pipeline_corr_producer_state); + ++pipeline_corr_producer_state; + + if constexpr (get<1>(ThreadShape{}) > 1) { + pipeline_kv.consumer_release(pipeline_kv_release_state); + ++pipeline_kv_release_state; + } + } + + // release Q1 + pipeline_q.consumer_release(pipeline_q_release_state); + ++pipeline_q_release_state; + + // release Q2 + if constexpr (get<0>(ThreadShape{}) > 1) { + pipeline_q.consumer_release(pipeline_q_release_state); + ++pipeline_q_release_state; + } + + // wait for Vi + if constexpr (get<1>(ThreadShape{}) > 1) { + v_index = pipeline_kv_consumer_state.index(); + pipeline_kv.consumer_wait(pipeline_kv_consumer_state); + ++pipeline_kv_consumer_state; + } + + // gemm P2 * Vi -> O2 + pipeline_corr.producer_acquire(pipeline_corr_producer_state); + pipeline_s1.producer_acquire(pipeline_s1_producer_state); + + gemm_reset_zero_acc(mma_pv_ts, tOrP1, tOrV(_, _, _, v_index), tOtO1); + + pipeline_corr.producer_commit(pipeline_corr_producer_state); + ++pipeline_corr_producer_state; + + // release Vi + pipeline_kv.consumer_release(pipeline_kv_release_state); + ++pipeline_kv_release_state; + + pipeline_s0.producer_commit(pipeline_s0_producer_state); + ++pipeline_s0_producer_state; + + pipeline_s1.producer_commit(pipeline_s1_producer_state); + ++pipeline_s1_producer_state; + + // T0 S00 B1, T0 S10 B1, T0 S00 B2, T0 S01 B1, T0 S10 B2, T0 S11 B1, T0 S01 B2, T1 S00 B1, T0 + // S11 B2, ... Q1 * K1 , Q2 * K1 , S11 * V1 , Q1 * K2 , S21 * V1 , Q2 * K2 , S12 * V2 , Q1 * + // K3 , S22 * K2 , ... + } + + template + CUTLASS_DEVICE auto softmax_step(float& row_max, float& row_sum, Stage stage, bool final_call, + BlkCoord const& blk_coord, CountingTensor const& cS, + Params const& params, ProblemShape const& problem_shape, + PipelineS& pipeline_s, + typename PipelineS::PipelineState& pipeline_s_consumer_state, + PipelineC& pipeline_c, + typename PipelineC::PipelineState& pipeline_c_producer_state, + OrderBarrierSoftmax& order_s) { + Tensor tScS = typename CollectiveMmaQK::TiledMma{}.get_slice(0).partition_C(cS); + + Tensor tStS = + partition_fragment_C(typename CollectiveMmaQK::TiledMma{}, select<0, 1>(TileShapeQK{})); + tStS.data() = uint32_t(stage == _0{} ? TmemAllocation::S0 : TmemAllocation::S1); + + Tensor tStS_v = tStS.compose(make_layout(make_shape(_128{}, _2{}))); + tStS_v.data() = uint32_t(stage == _0{} ? TmemAllocation::V0 : TmemAllocation::V1); + Tensor tScS_v = tScS.compose(make_layout(make_shape(_128{}, _2{}))); + + auto tilePlikeFP32 = get<1>(TileShapeQK{}) / Int{} * Int{}; + Tensor tStS_P = tStS.compose(make_layout(make_shape(_128{}, tilePlikeFP32))); + tStS_P.data() = warp_uniform(uint32_t(stage == _0{} ? TmemAllocation::P0 : TmemAllocation::P1)); + Tensor tScS_P = tScS.compose(make_layout(make_shape(_128{}, tilePlikeFP32))); + + // Each thread owns a single row + using TMEM_LOAD = SM100_TMEM_LOAD_32dp32b32x; // 4x32 threads with 128 cols of 32b elem + using TMEM_STORE = SM100_TMEM_STORE_32dp32b32x; // 4x32 threads with 128 cols of 8b elem + using TMEM_STORE_V = SM100_TMEM_STORE_32dp32b2x; // 4x32 threads with 2 cols of 32b elem + + int thread_idx = threadIdx.x % (4 * cutlass::NumThreadsPerWarp); + + auto tiled_tmem_load = make_tmem_copy(TMEM_LOAD{}, tStS); + auto thr_tmem_load = tiled_tmem_load.get_slice(thread_idx); + + Tensor tTMEM_LOADtS = thr_tmem_load.partition_S(tStS); + Tensor tTMEM_LOADcS = thr_tmem_load.partition_D(tScS); + + auto tiled_tmem_storev = make_tmem_copy(TMEM_STORE_V{}, tStS_v); + auto thr_tmem_storev = tiled_tmem_storev.get_slice(thread_idx); + + Tensor tTMEM_STOREVtS = thr_tmem_storev.partition_D(tStS_v); + Tensor tTMEM_STOREVcS = thr_tmem_storev.partition_S(tScS_v); + + auto tiled_tmem_store = make_tmem_copy(TMEM_STORE{}, tStS_P); + auto thr_tmem_store = tiled_tmem_store.get_slice(thread_idx); + + Tensor tTMEM_STOREtS_x4 = thr_tmem_store.partition_D(tStS_P); + tTMEM_STOREtS_x4.data() = warp_uniform(tTMEM_STOREtS_x4.data().get()); + Tensor tTMEM_STOREcS = thr_tmem_store.partition_S(tScS_P); + + // wait on tensor core pipe + pipeline_s.consumer_wait(pipeline_s_consumer_state); + + // read all of S from tmem into reg mem + Tensor tTMEM_LOADrS = make_tensor(shape(tTMEM_LOADcS)); + copy(tiled_tmem_load, tTMEM_LOADtS, tTMEM_LOADrS); + + if constexpr (need_apply_mask) { + Mask{}.apply_mask(tTMEM_LOADrS, tTMEM_LOADcS, problem_shape); + } + + ElementQK old_row_max = row_max; + { + // compute rowmax + float row_max_0 = row_max; + float row_max_1 = row_max; + float row_max_2 = row_max; + float row_max_3 = row_max; + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(tTMEM_LOADrS); i += 4) { + row_max_0 = ::fmax(row_max_0, tTMEM_LOADrS(i)); + row_max_1 = ::fmax(row_max_1, tTMEM_LOADrS(i + 1)); + row_max_2 = ::fmax(row_max_2, tTMEM_LOADrS(i + 2)); + row_max_3 = ::fmax(row_max_3, tTMEM_LOADrS(i + 3)); + } + row_max = ::fmax(row_max_0, row_max_1); + row_max = ::fmax(row_max, row_max_2); + row_max = ::fmax(row_max, row_max_3); + } + + ElementQK row_max_safe = row_max == -INFINITY ? 0 : row_max; + + Tensor tTMEM_STOREVrS = make_tensor(shape(tTMEM_STOREVcS)); + tTMEM_STOREVrS(kIdxOldRowMax) = old_row_max; + tTMEM_STOREVrS(kIdxNewRowMax) = row_max_safe; + copy(tiled_tmem_storev, tTMEM_STOREVrS, tTMEM_STOREVtS); + + pipeline_c.producer_commit(pipeline_c_producer_state); + ++pipeline_c_producer_state; + + // notify correction wg that they are ready (might need addtl ordering between S0 and S1 WG's) + + ElementQK scale = params.scale_softmax_log2; + ElementQK row_max_scale = row_max_safe * scale; + + float2 scale_fp32x2 = make_float2(scale, scale); + float2 minus_row_max_scale_fp32x2 = make_float2(-row_max_scale, -row_max_scale); + + Tensor tTMEM_STORErS_x4 = make_tensor(shape(tTMEM_STOREcS)); + + constexpr int kConversionsPerStep = 2; + + Tensor tTMEM_STORErS_x4_e = recast>(tTMEM_STORErS_x4); + + NumericArrayConverter convert; + + const int kReleasePipeCount = 10; // must be multiple of 2 + + order_s.wait(); + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(tTMEM_LOADrS); i += 2) { + float2 in = make_float2(tTMEM_LOADrS(i + 0), tTMEM_LOADrS(i + 1)); + float2 out; + cute::fma(out, scale_fp32x2, in, minus_row_max_scale_fp32x2); + tTMEM_LOADrS(i + 0) = out.x; + tTMEM_LOADrS(i + 1) = out.y; + + tTMEM_LOADrS(i + 0) = ::exp2f(tTMEM_LOADrS(i + 0)); + tTMEM_LOADrS(i + 1) = ::exp2f(tTMEM_LOADrS(i + 1)); + + Array in_conv; + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < kConversionsPerStep; j++) { + in_conv[j] = tTMEM_LOADrS(i + j); + } + tTMEM_STORErS_x4_e[i / kConversionsPerStep] = convert(in_conv); + + if (i == size(tTMEM_LOADrS) - kReleasePipeCount) { + order_s.arrive(); + } + + // this prevents register spills in fp16 + if constexpr (size<2>(tTMEM_STORErS_x4) == _2{}) { + if (i == size(tTMEM_LOADrS) - 6) { + copy(tiled_tmem_store, tTMEM_STORErS_x4(_, _, 0), tTMEM_STOREtS_x4(_, _, 0)); + } + } + } + + // tmem_store(reg_S8) -> op_P + CUTE_STATIC_ASSERT_V(size<2>(tTMEM_STORErS_x4) <= _2{}); + CUTE_STATIC_ASSERT_V(size<1>(tTMEM_STORErS_x4) == _1{}); + copy(tiled_tmem_store, tTMEM_STORErS_x4(_, _, size<2>(tTMEM_STORErS_x4) - 1), + tTMEM_STOREtS_x4(_, _, size<2>(tTMEM_STORErS_x4) - 1)); + + cutlass::arch::fence_view_async_tmem_store(); + + // notify tensor core warp that P is ready + pipeline_s.consumer_release(pipeline_s_consumer_state); + ++pipeline_s_consumer_state; + + pipeline_c.producer_acquire(pipeline_c_producer_state); + + ElementQK acc_scale = 0.5f * ::exp2f(scale * (old_row_max - row_max_safe)); + row_sum *= acc_scale; + // row_sum = sum(reg_S) + float2 local_row_sum_f32x2 = make_float2(row_sum, row_sum); + float2 local_row_sum_1 = make_float2(0, 0); + float2 local_row_sum_2 = make_float2(0, 0); + float2 local_row_sum_3 = make_float2(0, 0); + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(tTMEM_LOADrS); i += 8) { + // row_sum += tTMEM_LOADrS(i); + float2 in = make_float2(tTMEM_LOADrS(i), tTMEM_LOADrS(i + 1)); + cute::add(local_row_sum_f32x2, local_row_sum_f32x2, in); + + in = make_float2(tTMEM_LOADrS(i + 2), tTMEM_LOADrS(i + 2 + 1)); + cute::add(local_row_sum_1, local_row_sum_1, in); + + in = make_float2(tTMEM_LOADrS(i + 4), tTMEM_LOADrS(i + 4 + 1)); + cute::add(local_row_sum_2, local_row_sum_2, in); + + in = make_float2(tTMEM_LOADrS(i + 6), tTMEM_LOADrS(i + 6 + 1)); + cute::add(local_row_sum_3, local_row_sum_3, in); + } + + cute::add(local_row_sum_f32x2, local_row_sum_f32x2, local_row_sum_1); + cute::add(local_row_sum_2, local_row_sum_2, local_row_sum_3); + cute::add(local_row_sum_f32x2, local_row_sum_f32x2, local_row_sum_2); + float local_row_sum = local_row_sum_f32x2.x + local_row_sum_f32x2.y; + + row_sum = local_row_sum; + + if (final_call) { + // re-acquire the S part in the final step + pipeline_s.consumer_wait(pipeline_s_consumer_state); + + Tensor tTMEM_STOREVrS = make_tensor(shape(tTMEM_STOREVcS)); + tTMEM_STOREVrS(kIdxFinalRowMax) = row_max; + tTMEM_STOREVrS(kIdxFinalRowSum) = row_sum; + copy(tiled_tmem_storev, tTMEM_STOREVrS, tTMEM_STOREVtS); + } + } + + template + CUTLASS_DEVICE auto softmax(Stage stage, BlkCoord const& blk_coord, Params const& params, + ProblemShape const& problem_shape, PipelineS& pipeline_s, + typename PipelineS::PipelineState& pipeline_s_consumer_state, + PipelineC& pipeline_c, + typename PipelineC::PipelineState& pipeline_c_producer_state, + OrderBarrierSoftmax& order_s) { + int mask_tile_count = Mask{}.get_unmasked_trip_count(blk_coord, TileShape{}, problem_shape); + + ElementQK row_max = -INFINITY; + ElementQK row_sum = 0; + + Tensor cS_base = make_identity_tensor(select<0, 1>(TileShapeQK{})); + auto logical_offset = make_coord(get<0>(blk_coord) * get<0>(TileShape{}) + + (stage % get<0>(ThreadShape{})) * get<0>(TileShapeQK{}), + 0 + (stage % get<1>(ThreadShape{})) * get<1>(TileShapeQK{})); + Tensor cS = domain_offset(logical_offset, cS_base); + + pipeline_c.producer_acquire(pipeline_c_producer_state); + + CUTLASS_PRAGMA_NO_UNROLL + for (; mask_tile_count > 0; mask_tile_count -= 1) { + softmax_step( + row_max, row_sum, stage, + (mask_tile_count == 1) && + (Mask{}.get_masked_trip_count(blk_coord, TileShape{}, problem_shape) == 0), + blk_coord, cS, params, problem_shape, pipeline_s, pipeline_s_consumer_state, pipeline_c, + pipeline_c_producer_state, order_s); + + cS.data() = cS.data() + E<1>{} * get<1>(ThreadShape{}) * get<1>(TileShapeQK{}); + } + + // Masked iterations + mask_tile_count = Mask{}.get_masked_trip_count(blk_coord, TileShape{}, problem_shape); + + CUTLASS_PRAGMA_NO_UNROLL + for (; mask_tile_count > 0; mask_tile_count -= 1) { + softmax_step( + row_max, row_sum, stage, mask_tile_count == 1, blk_coord, cS, params, problem_shape, + pipeline_s, pipeline_s_consumer_state, pipeline_c, pipeline_c_producer_state, order_s); + + cS.data() = cS.data() + E<1>{} * get<1>(ThreadShape{}) * get<1>(TileShapeQK{}); + } + + pipeline_c.producer_commit(pipeline_c_producer_state); + ++pipeline_c_producer_state; + + pipeline_c.producer_acquire(pipeline_c_producer_state); + // empty step to sync against pipe s + pipeline_s.consumer_release(pipeline_s_consumer_state); + ++pipeline_s_consumer_state; + } + + template + CUTLASS_DEVICE auto correction_epilogue(float scale_softmax_log2, float scale_out, + Vector const& v0, Vector const& v1, GTensor& gO, + CTensor const& cO, Shape const& g_shape, + Epilogue const& epilogue) { + using ElementOut = typename GTensor::value_type; + + int thread_idx = threadIdx.x % (4 * cutlass::NumThreadsPerWarp); + + // As opposed to the softmax, we do not have enough registers here + // to load all of the values (for tile kv = 128), so we loop + // good values would be either 32 or 64 + const int kCorrectionTileSize = 32 / sizeof(ElementOut); + + using TMEM_LOAD = + std::conditional_t; // 4x32 threads with 64 cols of 32b elem + + typename CollectiveMmaPV::TiledMma mma; + Tensor tOtO = partition_fragment_C(mma, select<0, 1>(TileShapePV{})); + Tensor tOcO = mma.get_slice(0).partition_C(cO); + Tensor tOgO = mma.get_slice(0).partition_C(gO); + + Tensor tOtO_i = tOtO.compose(make_layout(make_shape(_128{}, Int{}))); + Tensor tOcO_i = tOcO.compose(make_layout(make_shape(_128{}, Int{}))); + Tensor tOgO_i = tOgO.compose(make_layout(make_shape(_128{}, Int{}))); + + Tensor tOtO0 = tOtO_i; + tOtO0.data() = tOtO0.data().get() + uint32_t(TmemAllocation::O0); + Tensor tOtO1 = tOtO_i; + tOtO1.data() = tOtO1.data().get() + uint32_t(TmemAllocation::O1); + + auto tiled_tmem_load = make_tmem_copy(TMEM_LOAD{}, tOtO_i); + auto thr_tmem_load = tiled_tmem_load.get_slice(thread_idx); + + Tensor tTMEM_LOADtO0 = thr_tmem_load.partition_S(tOtO0); + Tensor tTMEM_LOADtO1 = thr_tmem_load.partition_S(tOtO1); + Tensor tTMEM_LOADcO = thr_tmem_load.partition_D(tOcO_i); + Tensor tTMEM_LOADgO = thr_tmem_load.partition_D(tOgO_i); + + float row_max = std::max(v0(kIdxFinalRowMax), v1(kIdxFinalRowMax)); + float adj0 = ::exp2f(scale_softmax_log2 * (v0(kIdxFinalRowMax) - row_max)); + float adj1 = ::exp2f(scale_softmax_log2 * (v1(kIdxFinalRowMax) - row_max)); + float row_sum = adj0 * v0(kIdxFinalRowSum) + adj1 * v1(kIdxFinalRowSum); + float scale0 = scale_out * adj0 / row_sum; + float scale1 = scale_out * adj1 / row_sum; + + float2 scale0_f32x2 = make_float2(scale0, scale0); + float2 scale1_f32x2 = make_float2(scale1, scale1); + + // loop: + // TMEM_LOAD, TMEM_LOAD, FMUL2, FFMA2, STG + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < 128 / kCorrectionTileSize; i++) { + Tensor tTMEM_LOADtO0_i = tTMEM_LOADtO0; + tTMEM_LOADtO0_i.data() = tTMEM_LOADtO0_i.data().get() + uint32_t(i * kCorrectionTileSize); + Tensor tTMEM_LOADtO1_i = tTMEM_LOADtO1; + tTMEM_LOADtO1_i.data() = tTMEM_LOADtO1_i.data().get() + uint32_t(i * kCorrectionTileSize); + Tensor tTMEM_LOADgO_i = tTMEM_LOADgO; + tTMEM_LOADgO_i.data() = tTMEM_LOADgO_i.data().get() + i * kCorrectionTileSize * stride<1>(gO); + + Tensor tTMrO0 = make_tensor(shape(tTMEM_LOADcO)); + Tensor tTMrO1 = make_tensor(shape(tTMEM_LOADcO)); + + copy(tiled_tmem_load, tTMEM_LOADtO0_i, tTMrO0); + copy(tiled_tmem_load, tTMEM_LOADtO1_i, tTMrO1); + + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < size(tTMrO0); j += 2) { + float2 in0 = make_float2(tTMrO0(j), tTMrO0(j + 1)); + float2 in1 = make_float2(tTMrO1(j), tTMrO1(j + 1)); + float2 out; + cute::mul(out, scale0_f32x2, in0); + cute::fma(out, scale1_f32x2, in1, out); + tTMrO0(j) = out.x; + tTMrO0(j + 1) = out.y; + } + + constexpr int N = 4 / sizeof(ElementOut); + NumericArrayConverter convert; + + Tensor tSMrO = make_tensor_like(tTMrO0); + + Tensor tCs = recast(tTMrO0); + Tensor tCd = recast(tSMrO); + + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < size(tCs); j++) { + tCd(j) = convert.convert(tCs(j)); + } + + Tensor tSMgO_i = recast(tTMEM_LOADgO_i); + Tensor tSMrO_i = recast(tSMrO); + + // could use masking do this right for smaller D + if (get<0>(tTMEM_LOADcO(_0{})) < get<0>(g_shape)) { + copy(AutoVectorizingCopyWithAssumedAlignment<128>{}, tSMrO_i, tSMgO_i); + } + } + } + + CUTLASS_DEVICE auto correction_rescale(float scale, uint32_t tmem_O) { + int thread_idx = threadIdx.x % (4 * cutlass::NumThreadsPerWarp); + + // As opposed to the softmax, we do not have enough registers here + // to load all of the values (for tile kv = 128), so we loop + // good values would be either 32 or 64 + const int kCorrectionTileSize = 32; + + using TMEM_LOAD = SM100_TMEM_LOAD_32dp32b32x; // 4x32 threads with 64 cols of 32b elem + using TMEM_STORE = SM100_TMEM_STORE_32dp32b32x; // 4x32 threads with 64 cols of 32b elem + + typename CollectiveMmaPV::TiledMma mma; + Tensor cO = make_identity_tensor(select<0, 1>(TileShapePV{})); + Tensor tOtO = partition_fragment_C(mma, select<0, 1>(TileShapePV{})); + Tensor tOcO = mma.get_slice(0).partition_C(cO); + + Tensor tOtO_i = tOtO.compose(make_layout(make_shape(_128{}, Int{}))); + Tensor tOcO_i = tOcO.compose(make_layout(make_shape(_128{}, Int{}))); + + tOtO_i.data() = tOtO_i.data().get() + tmem_O; + + auto tiled_tmem_load = make_tmem_copy(TMEM_LOAD{}, tOtO_i); + auto thr_tmem_load = tiled_tmem_load.get_slice(thread_idx); + auto tiled_tmem_store = make_tmem_copy(TMEM_STORE{}, tOtO_i); + auto thr_tmem_store = tiled_tmem_store.get_slice(thread_idx); + + Tensor tTMEM_LOADtO = thr_tmem_load.partition_S(tOtO_i); + Tensor tTMEM_LOADcO = thr_tmem_load.partition_D(tOcO_i); + Tensor tTMEM_STOREtO = thr_tmem_store.partition_D(tOtO_i); + Tensor tTMEM_STOREcO = thr_tmem_store.partition_S(tOcO_i); + static_assert(shape(tTMEM_STOREcO) == shape(tTMEM_LOADcO)); + + float2 scale_f32x2 = make_float2(scale, scale); + + Tensor tTMrO = + make_tensor(make_shape(shape(tTMEM_LOADcO), Int<128 / kCorrectionTileSize>{})); + + auto copy_in = [&](int i) { + Tensor tTMEM_LOADtO_i = tTMEM_LOADtO; + tTMEM_LOADtO_i.data() = tTMEM_LOADtO_i.data().get() + uint32_t(i * kCorrectionTileSize); + Tensor tTMrO_i = tTMrO(_, i).compose(make_layout(shape<0>(tTMrO))); + copy(tiled_tmem_load, tTMEM_LOADtO_i, tTMrO_i); + }; + + auto copy_out = [&](int i) { + Tensor tTMEM_STOREtO_i = tTMEM_STOREtO; + tTMEM_STOREtO_i.data() = tTMEM_STOREtO_i.data().get() + uint32_t(i * kCorrectionTileSize); + Tensor tTMrO_i = tTMrO(_, i).compose(make_layout(shape<0>(tTMrO))); + copy(tiled_tmem_store, tTMrO_i, tTMEM_STOREtO_i); + }; + + // sequence: LLMSLMSLMSS + + // loop: + // TMEM_LOAD, FMUL2 scale, TMEM_STORE + copy_in(0); + + int count = get<2>(TileShape{}) / kCorrectionTileSize; + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < count; i++) { + if (i != count - 1) { + copy_in(i + 1); + } + + Tensor tTMrO_i = tTMrO(_, i).compose(make_layout(shape<0>(tTMrO))); + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < size(tTMrO_i); j += 2) { + float2 in = make_float2(tTMrO_i(j), tTMrO_i(j + 1)); + float2 out; + cute::mul(out, scale_f32x2, in); + tTMrO_i(j) = out.x; + tTMrO_i(j + 1) = out.y; + } + + copy_out(i); + } + } + + template + CUTLASS_DEVICE auto correction( + BlkCoord const& blk_coord, Params const& params, ProblemShape const& problem_shape, + TensorStorageEpi& shared_storage_epi, PipelineC& pipeline_s0_c, + typename PipelineC::PipelineState& pipeline_s0_c_consumer_state, PipelineC& pipeline_s1_c, + typename PipelineC::PipelineState& pipeline_s1_c_consumer_state, PipelineO& pipeline_o, + typename PipelineO::PipelineState& pipeline_o_consumer_state, PipelineE& pipeline_epi, + typename PipelineE::PipelineState& pipeline_epi_producer_state, Epilogue const& epilogue) { + int mask_tile_count = Mask{}.get_trip_count(blk_coord, TileShape{}, problem_shape); + + int thread_idx = threadIdx.x % (4 * cutlass::NumThreadsPerWarp); + + Tensor tStS = + partition_fragment_C(typename CollectiveMmaQK::TiledMma{}, select<0, 1>(TileShapeQK{})); + + Tensor cS = make_identity_tensor(select<0, 1>(TileShapeQK{})); + Tensor tScS = typename CollectiveMmaQK::TiledMma{}.get_slice(0).partition_C(cS); + + Tensor tStS_v = tStS.compose(make_layout(make_shape(_128{}, _2{}))); + Tensor tScS_v = tScS.compose(make_layout(make_shape(_128{}, _2{}))); + + using TMEM_LOAD_V = SM100_TMEM_LOAD_32dp32b2x; // 4x32 threads with 2 cols of 32b elem + + auto tiled_tmem_loadv = make_tmem_copy(TMEM_LOAD_V{}, tStS_v); + auto thr_tmem_loadv = tiled_tmem_loadv.get_slice(thread_idx); + + Tensor tTMEM_LOADVtS = thr_tmem_loadv.partition_S(tStS_v); + Tensor tTMEM_LOADVcS = thr_tmem_loadv.partition_D(tScS_v); + + Tensor tTMEM_LOADVtS0 = tTMEM_LOADVtS; + tTMEM_LOADVtS0.data() = tTMEM_LOADVtS0.data().get() + uint32_t(TmemAllocation::V0); + Tensor tTMEM_LOADVtS1 = tTMEM_LOADVtS; + tTMEM_LOADVtS1.data() = tTMEM_LOADVtS1.data().get() + uint32_t(TmemAllocation::V1); + + // ignore first signal from softmax as no correction is required + pipeline_s0_c.consumer_wait(pipeline_s0_c_consumer_state); + pipeline_s0_c.consumer_release(pipeline_s0_c_consumer_state); + ++pipeline_s0_c_consumer_state; + + pipeline_s1_c.consumer_wait(pipeline_s1_c_consumer_state); + + // handle the last iteration differently (i.e. tmem_load/stsm for epi) + mask_tile_count -= 1; + + CUTLASS_PRAGMA_NO_UNROLL + for (; mask_tile_count > 0; mask_tile_count -= 1) { + pipeline_s0_c.consumer_wait(pipeline_s0_c_consumer_state); + + Tensor tTMEM_LOADVrS = make_tensor(shape(tTMEM_LOADVcS)); + + // read row_wise new global max + copy(tiled_tmem_loadv, tTMEM_LOADVtS0, tTMEM_LOADVrS); + + // e^(scale * (old_max - new_max) + float scale = ::exp2f(params.scale_softmax_log2 * + (tTMEM_LOADVrS(kIdxOldRowMax) - tTMEM_LOADVrS(kIdxNewRowMax))); + + pipeline_o.consumer_wait(pipeline_o_consumer_state); + + correction_rescale(scale, uint32_t(TmemAllocation::O0)); + + pipeline_s1_c.consumer_release(pipeline_s1_c_consumer_state); + ++pipeline_s1_c_consumer_state; + + cutlass::arch::fence_view_async_tmem_store(); + + pipeline_o.consumer_release(pipeline_o_consumer_state); + ++pipeline_o_consumer_state; + + pipeline_s1_c.consumer_wait(pipeline_s1_c_consumer_state); + + copy(tiled_tmem_loadv, tTMEM_LOADVtS1, tTMEM_LOADVrS); + + scale = ::exp2f(params.scale_softmax_log2 * + (tTMEM_LOADVrS(kIdxOldRowMax) - tTMEM_LOADVrS(kIdxNewRowMax))); + + pipeline_o.consumer_wait(pipeline_o_consumer_state); + + correction_rescale(scale, uint32_t(TmemAllocation::O1)); + + pipeline_s0_c.consumer_release(pipeline_s0_c_consumer_state); + ++pipeline_s0_c_consumer_state; + + cutlass::arch::fence_view_async_tmem_store(); + + pipeline_o.consumer_release(pipeline_o_consumer_state); + ++pipeline_o_consumer_state; + } + + pipeline_s1_c.consumer_release(pipeline_s1_c_consumer_state); + ++pipeline_s1_c_consumer_state; + + // do the final correction to O1 + // better to somehow special-case it in the loop above + // doesn't matter for non-persistent code, but if it were + // persistent we do not want to release O too early + + pipeline_s0_c.consumer_wait(pipeline_s0_c_consumer_state); + + // read from V0 + // read row_sum and final row_max here + Tensor tTMEM_LOADVrS0 = make_tensor(shape(tTMEM_LOADVcS)); + copy(tiled_tmem_loadv, tTMEM_LOADVtS0, tTMEM_LOADVrS0); + + pipeline_s0_c.consumer_release(pipeline_s0_c_consumer_state); + ++pipeline_s0_c_consumer_state; + + pipeline_s1_c.consumer_wait(pipeline_s1_c_consumer_state); + + // load from V1 + Tensor tTMEM_LOADVrS1 = make_tensor(shape(tTMEM_LOADVcS)); + copy(tiled_tmem_loadv, tTMEM_LOADVtS1, tTMEM_LOADVrS1); + + pipeline_s1_c.consumer_release(pipeline_s1_c_consumer_state); + ++pipeline_s1_c_consumer_state; + + auto pipeline_o_release_state = pipeline_o_consumer_state; + pipeline_o.consumer_wait(pipeline_o_consumer_state); + ++pipeline_o_consumer_state; + pipeline_o.consumer_wait(pipeline_o_consumer_state); + ++pipeline_o_consumer_state; + // store to epi smem + + // loop: + // TMEM_LOAD + // FMUL2 scale = 1 / global_sum * out_quant_scale + // F2FP + // store to smem + + Tensor cO = make_identity_tensor(select<0, 1>(TileShapePV{})); + auto g_shape = select<0, 2>(problem_shape); + auto mO = make_tensor(make_gmem_ptr(epilogue.params.ptr_o), + append<3>(select<0, 1>(TileShapePV{}), get<3>(problem_shape)), + epilogue.params.dO); + auto gO = mO(_, _, get<2>(blk_coord)); + + correction_epilogue(params.scale_softmax_log2, params.scale_output, tTMEM_LOADVrS0, + tTMEM_LOADVrS1, gO, cO, g_shape, epilogue); + + cutlass::arch::fence_view_async_tmem_load(); + + pipeline_o.consumer_release(pipeline_o_release_state); + ++pipeline_o_release_state; + + pipeline_o.consumer_release(pipeline_o_release_state); + ++pipeline_o_release_state; + } +}; + +} // namespace cutlass::fmha::collective diff --git a/include/flashinfer/attention/blackwell/collective/sm100_fmha_load_cpasync_warpspecialized.hpp b/include/flashinfer/attention/blackwell/collective/sm100_fmha_load_cpasync_warpspecialized.hpp new file mode 100644 index 0000000000000000000000000000000000000000..0d35a60b834b36f02baf653b6a5063f8bca87b92 --- /dev/null +++ b/include/flashinfer/attention/blackwell/collective/sm100_fmha_load_cpasync_warpspecialized.hpp @@ -0,0 +1,362 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include "collective/fmha_common.hpp" +#include "collective/fmha_fusion.hpp" +#include "cute/layout.hpp" +#include "cute/tensor.hpp" +#include "cutlass/arch/memory_sm80.h" +#include "cutlass/cutlass.h" +#include "cutlass/gemm/collective/collective_builder.hpp" + +namespace cutlass::fmha::collective { + +using namespace cute; + +template +struct Sm100FmhaLoadCpAsyncWarpspecialized { + using TileShapeQK = typename CollectiveMmaQK::TileShape; + using TileShapePV = typename CollectiveMmaPV::TileShape; + + struct Arguments { + const int* cache_batch_idx; + + const Element* ptr_q; + StrideQ dQ; + + const Element* ptr_new_k; + StrideNewK dNewK; + const Element* ptr_new_v; + StrideNewV dNewV; + + Element* ptr_cache_k; + StrideCacheK dCacheK; + Element* ptr_cache_v; + StrideCacheV dCacheV; + }; + + using Params = Arguments; + + template + static Params to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, + void* workspace) { + return args; + } + + CUTLASS_DEVICE + static void prefetch_tma_descriptors(Params const& params) {} + + template + CUTLASS_DEVICE auto constexpr transpose(Tensor const& t) { + CUTE_STATIC_ASSERT_V(rank(t) == _2{}); + return t.compose( + make_layout(make_shape(size<1>(t), size<0>(t)), make_stride(size<0>(t), _1{}))); + } + + template + CUTLASS_DEVICE void copy_with_limit(TiledCopy const& tiled_copy, + CountTensor const& c, CountLimit const& l, + SrcTensor const& src, DstTensor&& dst) { + // copy(tiled_copy, src, dst); +#if 1 + auto c_f = make_tensor(c.data(), flatten(c.layout())); + auto src_f = make_tensor(src.data(), flatten(src.layout())); + auto dst_f = make_tensor(dst.data(), flatten(dst.layout())); + auto c_v = group_modes<1, rank_v>(c_f); + auto src_v = group_modes<1, rank_v>(src_f); + auto dst_v = group_modes<1, rank_v>(dst_f); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size<1>(src_v); i++) { + if (elem_less(c_v(_0{}, i), l)) { + copy(CAtom{}, src_v(_, i), dst_v(_, i)); + } else { + clear(dst_v(_, i)); + } + } +#endif + } + + template + CUTLASS_DEVICE void load(BlkCoord const& blk_coord, ProblemShape const& problem_shape, + Params const& params, ParamsProblemShape const& params_problem_shape, + TensorStorage& storage, PipelineQ& pipeline_q, + typename PipelineQ::PipelineState& pipeline_q_producer_state, + PipelineKV& pipeline_kv, + typename PipelineKV::PipelineState& pipeline_kv_producer_state) { + int mask_tile_count = Mask{}.get_trip_count(blk_coord, TileShape{}, problem_shape); + mask_tile_count *= 2; + + int warp_idx = (threadIdx.x / 32) % 2; + int thread_idx = warp_idx * 32 + (threadIdx.x % 32); + + using X = Underscore; + + // this one is only executed by one thread, no need to elect_one + auto blk_coord_cache = blk_coord; + if (params.cache_batch_idx != nullptr) { + get<2, 1>(blk_coord_cache) = params.cache_batch_idx[get<2, 1>(blk_coord_cache)]; + } + + // Q1, K1, K2, V1, K3, V2, ... Kn, Vn-1, Vn + // two pipes: Q and KV + auto cQ = make_identity_tensor(select<0, 2>(TileShape{})); + auto mQ = make_tensor(make_gmem_ptr(params.ptr_q), + append<3>(select<0, 2>(TileShapeQK{}), get<3>(problem_shape)), params.dQ); + auto gQ = mQ(_, _, get<2>(blk_coord)); + auto sQ = make_tensor(make_smem_ptr(storage.smem_q.data()), SmemLayoutQ{}); + + typename CollectiveMmaQK::TiledMma mma_qk; + ThrMMA thr_mma_qk = mma_qk.get_slice(0); + auto tSgQ = thr_mma_qk.partition_A(gQ); + auto tScQ = thr_mma_qk.partition_A(cQ); + + auto atom_q_tv = Layout, Shape<_16, _16>>, + Stride, Stride<_1, _1024>>>{}; + auto atom_kv_tv = Layout, Shape<_16, _4>>, + Stride, Stride<_1, _1024>>>{}; + + auto tiled_copy_q = make_cotiled_copy( + Copy_Atom, Element>{}, atom_q_tv, + make_layout(shape(tSgQ), + replace<0>(stride(tSgQ), replace<0>(stride<0>(tSgQ), get<2>(TileShape{}))))); + + auto thr_copy_q = tiled_copy_q.get_slice(thread_idx); + + auto tQsQ = thr_copy_q.partition_D(sQ); + auto tQgQ = thr_copy_q.partition_S(tSgQ); + auto tQcQ = thr_copy_q.partition_S(tScQ); + + auto limitQ = append<2>(get<0>(problem_shape), _128{}); + + // Q1 + int q0_index = get<0>(blk_coord); + + auto load_q = [&](int q_index, auto& state) { + pipeline_q.producer_acquire(state); + + // q is always loaded masked + using Vec = uint128_t; + Vec vzero = uint128_t(0, 0); + auto src = recast(tQgQ(_, _, _, _)); + auto dst = recast(tQsQ(_, _, _, _, state.index())); + auto c = tQcQ(_, _, _, _); + int vlen = sizeof(Vec) / sizeof(Element); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(src); i++) { + auto cc = c(vlen * i); + Vec* dst_ptr = &dst(i); + const Vec* src_ptr = &src(i); + bool guard = elem_less(cc, limitQ); + cutlass::arch::cp_async_zfill<16, cutlass::arch::CacheOperation::Always>(dst_ptr, src_ptr, + guard); + } + + pipeline_q.producer_commit(state, cutlass::arch::cpasync_barrier_arrive); + }; + + load_q(q0_index, pipeline_q_producer_state); + ++pipeline_q_producer_state; + + auto cK_t = make_identity_tensor(select<1, 2>(TileShapeQK{})); + auto cK = make_tensor(cK_t.data(), + make_layout(get<0>(cK_t.layout()), get<1>(cK_t.layout()), + make_layout(_2{}, get<1>(TileShapeQK{}) * stride<0>(cK_t)))); + auto mK = make_tensor(make_gmem_ptr(params.ptr_cache_k), select<1, 2, 3>(problem_shape), + params.dCacheK); + auto gK = local_tile(mK(_, _, get<2>(blk_coord_cache)), TileShapeQK{}, make_coord(_, _, _0{}), + Step{}); + auto sK = make_tensor(make_smem_ptr(storage.smem_k.data()), SmemLayoutK{}); + + auto tSgK = thr_mma_qk.partition_B(gK); + auto tScK = thr_mma_qk.partition_B(cK); + + auto tSlK = thr_mma_qk.partition_B(make_tensor( + (Element*)nullptr, make_ordered_layout(select<1, 2>(TileShapeQK{}), Step<_1, _0>{}))); + auto tiled_copy_k = make_cotiled_copy( + Copy_Atom, Element>{}, atom_kv_tv, tSlK.layout()); + + auto thr_copy_k = tiled_copy_k.get_slice(thread_idx); + + auto tKsK = thr_copy_k.partition_D(sK); + auto tKgK = thr_copy_k.partition_S(tSgK); + auto tKcK = thr_copy_k.partition_S(tScK); + + int seqlen_cache_kv = get<1>(problem_shape) - ((params.ptr_new_k != nullptr) ? 1 : 0); + auto limitK = append<2>(seqlen_cache_kv, _128{}); + + auto cV_t = make_identity_tensor(select<1, 2>(TileShapePV{})); + auto cV = make_tensor(cV_t.data(), + make_layout(get<0>(cV_t.layout()), get<1>(cV_t.layout()), + make_layout(_2{}, get<2>(TileShapePV{}) * stride<1>(cV_t)))); + auto mV = make_tensor(make_gmem_ptr(params.ptr_cache_v), select<2, 1, 3>(problem_shape), + select<1, 0, 2>(params.dCacheV)); + auto gV = local_tile(mV(_, _, get<2>(blk_coord_cache)), TileShapePV{}, make_coord(_, _0{}, _), + Step{}); + auto sV = make_tensor(make_smem_ptr(storage.smem_v.data()), SmemLayoutV{}); + + typename CollectiveMmaPV::TiledMma mma_pv; + ThrMMA thr_mma_pv = mma_pv.get_slice(0); + auto tOgV = thr_mma_pv.partition_B(gV); + auto tOcV = thr_mma_pv.partition_B(cV); + auto tOlV = thr_mma_pv.partition_B( + make_tensor((Element*)nullptr, make_layout(select<1, 2>(TileShapePV{})))); + + auto tiled_copy_v = make_cotiled_copy( + Copy_Atom, Element>{}, atom_kv_tv, tOlV.layout()); + + auto thr_copy_v = tiled_copy_v.get_slice(thread_idx); + + auto tVsV = thr_copy_v.partition_D(sV); + auto tVgV = thr_copy_v.partition_S(tOgV); + auto tVcV = thr_copy_v.partition_S(tOcV); + + auto limitV = select<1, 0>(limitK); + + int full_tiles_cache = seqlen_cache_kv / get<1>(TileShapeQK{}); + + bool has_new = params.ptr_new_k != nullptr; + Tensor mNewK = + make_tensor(make_gmem_ptr(params.ptr_new_k), select<1, 2, 3>(problem_shape), params.dNewK); + Tensor mNewV = + make_tensor(make_gmem_ptr(params.ptr_new_v), select<1, 2, 3>(problem_shape), params.dNewV); + Tensor gNewK = mNewK(_, _, get<2>(blk_coord)); + Tensor gNewV = mNewV(_, _, get<2>(blk_coord)); + + auto load_k = [&](int k_index, auto& state) { + pipeline_kv.producer_acquire(state); + + if (k_index < full_tiles_cache) { + copy(tiled_copy_k, tKgK(_, _, _, _, k_index), tKsK(_, _, _, _, state.index())); + pipeline_kv.producer_commit(state, cutlass::arch::cpasync_barrier_arrive); + } else { + using Vec = uint128_t; + Vec vzero = uint128_t(0, 0); + auto src = recast(tKgK(_, _, _, _, k_index)); + auto dst = recast(tKsK(_, _, _, _, state.index())); + auto src2 = recast(gNewK); + auto c = tKcK(_, _, _, _, k_index); + int vlen = sizeof(Vec) / sizeof(Element); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(src); i++) { + auto cc = c(vlen * i); + Vec* dst_ptr = &dst(i); + const Vec* src_ptr = &src(i); + bool guard = elem_less(cc, limitK); + if (get<0>(cc) == seqlen_cache_kv && has_new) { + src_ptr = &src2(_0{}, get<1>(cc) / vlen); + guard = true; + } + cutlass::arch::cp_async_zfill<16, cutlass::arch::CacheOperation::Global>(dst_ptr, src_ptr, + guard); + } + + pipeline_kv.producer_commit(state, cutlass::arch::cpasync_barrier_arrive); + } + }; + + auto load_v = [&](int v_index, auto& state) { + pipeline_kv.producer_acquire(state); + + if (v_index < full_tiles_cache) { + copy(tiled_copy_v, tVgV(_, _, _, _, v_index), tVsV(_, _, _, _, state.index())); + pipeline_kv.producer_commit(state, cutlass::arch::cpasync_barrier_arrive); + } else { + using Vec = uint128_t; + Vec vzero = uint128_t(0, 0); + auto src = recast(tVgV(_, _, _, _, v_index)); + auto dst = recast(tVsV(_, _, _, _, state.index())); + auto src2 = recast(gNewV); + int vlen = sizeof(Vec) / sizeof(Element); + auto c = tVcV(_, _, _, _, v_index); + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(src); i++) { + auto cc = c(vlen * i); + Vec* dst_ptr = &dst(i); + const Vec* src_ptr = &src(i); + bool guard = elem_less(cc, limitV); + if (get<1>(cc) == seqlen_cache_kv && has_new) { + src_ptr = &src2(_0{}, get<0>(cc) / vlen); + guard = true; + } + cutlass::arch::cp_async_zfill<16, cutlass::arch::CacheOperation::Global>(dst_ptr, src_ptr, + guard); + } + + pipeline_kv.producer_commit(state, cutlass::arch::cpasync_barrier_arrive); + } + }; + + // K1 + int k_index = 0; + int v_index = 0; + + load_k(k_index, pipeline_kv_producer_state); + + ++pipeline_kv_producer_state; + k_index += 1; + + mask_tile_count -= 1; + + for (; mask_tile_count > 0; mask_tile_count -= 1) { + load_k(k_index, pipeline_kv_producer_state); + + ++pipeline_kv_producer_state; + k_index += 1; + + load_v(v_index, pipeline_kv_producer_state); + + ++pipeline_kv_producer_state; + v_index += 1; + } + + // V1 + + load_v(v_index, pipeline_kv_producer_state); + + ++pipeline_kv_producer_state; + v_index += 1; + + if (has_new) { + for (int i = thread_idx; i < get<2>(TileShape{}); i += 64) { + gK(seqlen_cache_kv, i, 0) = gNewK(0, i); + gV(i, seqlen_cache_kv, 0) = gNewV(0, i); + } + } + } +}; + +} // namespace cutlass::fmha::collective diff --git a/include/flashinfer/attention/blackwell/collective/sm100_fmha_load_tma_warpspecialized.hpp b/include/flashinfer/attention/blackwell/collective/sm100_fmha_load_tma_warpspecialized.hpp new file mode 100644 index 0000000000000000000000000000000000000000..4524b61b4c8da03a47da19c7fd56971535042740 --- /dev/null +++ b/include/flashinfer/attention/blackwell/collective/sm100_fmha_load_tma_warpspecialized.hpp @@ -0,0 +1,265 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include "../../../cutlass_utils.cuh" +#include "cute/layout.hpp" +#include "cute/tensor.hpp" +#include "cutlass/arch/memory_sm80.h" +#include "cutlass/cutlass.h" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "fmha_common.hpp" +#include "fmha_fusion.hpp" + +namespace cutlass::fmha::collective { + +using namespace cute; + +template +struct Sm100FmhaLoadTmaWarpspecialized { + using TileShapeQK = typename CollectiveMmaQK::TileShape; + using TileShapePV = typename CollectiveMmaPV::TileShape; + + using GmemTiledCopyQ = cute::SM90_TMA_LOAD; + using GmemTiledCopyKV = cute::SM90_TMA_LOAD; + static constexpr uint32_t NumStagesQ = PipelineQ::Stages; + + // (N, D, (H_R, H_G)) + using ShapeT = cute::Shape>; + // (N, D, (H_R, H_G)) + using StrideQ = cute::Shape>; + using StrideK = cute::Shape>; + using StrideV = cute::Shape<_1, int32_t, cute::Shape<_0, int32_t>>; + using LayoutQ = cute::Layout; + using LayoutK = cute::Layout; + using LayoutV = cute::Layout; + struct Arguments { + const Element* ptr_Q; + LayoutQ layout_Q; + const Element* ptr_K; + LayoutK layout_K; + const Element* ptr_V; + LayoutV layout_V; + }; + + // using ShapeLseT = cute::Shape; + // using StrideLseT = cute::Shape<_1, int64_t>; + // using LayoutLseT = cute::Layout; + + using ClusterLayout_VMNK = + decltype(tiled_divide(make_layout(Shape<_1, _1, _1>{}), + make_tile(typename CollectiveMmaQK::TiledMma::AtomThrID{}))); + using TMA_Q = typename CollectiveMmaQK::Params::TMA_A; + using TMA_K = typename CollectiveMmaQK::Params::TMA_B; + using TMA_V = typename CollectiveMmaPV::Params::TMA_B; + + struct Params { + TMA_Q tma_load_Q; + LayoutQ layout_Q; + TMA_K tma_load_K; + LayoutK layout_K; + TMA_V tma_load_V; + LayoutV layout_V; + }; + + template + static Params to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, + void* workspace) { + static_assert(is_variable_length_v>); + static_assert(is_variable_length_v>); + auto ptr_Q = args.ptr_Q; + auto ptr_K = args.ptr_K; + auto ptr_V = args.ptr_V; + LayoutQ layout_Q = args.layout_Q; + LayoutK layout_K = args.layout_K; + LayoutV layout_V = args.layout_V; + + auto mQ = make_tensor(make_gmem_ptr(ptr_Q), layout_Q); + auto mK = make_tensor(make_gmem_ptr(ptr_K), layout_K); + auto mV = make_tensor(make_gmem_ptr(ptr_V), layout_V); + + auto cluster_layout_vmnk = + tiled_divide(make_layout(Shape<_1, _1, _1>{}), + make_tile(typename CollectiveMmaQK::TiledMma::AtomThrID{})); + TMA_Q tma_load_Q = make_tma_atom_A_sm100( + GmemTiledCopyQ{}, mQ, SmemLayoutQ{}(_, _, _, _0{}), TileShapeQK{}, + typename CollectiveMmaQK::TiledMma{}, cluster_layout_vmnk); + TMA_K tma_load_K = make_tma_atom_B_sm100( + GmemTiledCopyKV{}, mK, SmemLayoutK{}(_, _, _, _0{}), TileShapeQK{}, + typename CollectiveMmaQK::TiledMma{}, cluster_layout_vmnk); + TMA_V tma_load_V = make_tma_atom_B_sm100( + GmemTiledCopyKV{}, mV, SmemLayoutV{}(_, _, _, _0{}), TileShapePV{}, + typename CollectiveMmaPV::TiledMma{}, cluster_layout_vmnk); + + return Params{tma_load_Q, layout_Q, tma_load_K, layout_K, tma_load_V, layout_V}; + } + + CUTLASS_DEVICE + static void prefetch_tma_descriptors(Params const& params) { + cute::prefetch_tma_descriptor(params.tma_load_Q.get_tma_descriptor()); + cute::prefetch_tma_descriptor(params.tma_load_K.get_tma_descriptor()); + cute::prefetch_tma_descriptor(params.tma_load_V.get_tma_descriptor()); + } + + template + CUTLASS_DEVICE void load(BlkCoord const& blk_coord, ProblemShape const& problem_shape, + Params const& params, ParamsProblemShape const& params_problem_shape, + TensorStorage& storage, PipelineQ& pipeline_q, + typename PipelineQ::PipelineState& pipeline_q_producer_state, + PipelineK& pipeline_k, + typename PipelineK::PipelineState& pipeline_k_producer_state, + PipelineV& pipeline_v, + typename PipelineV::PipelineState& pipeline_v_producer_state) { + int qo_tile_idx = get<0>(blk_coord); + int qo_head_idx = get<2, 0>(blk_coord); + int batch_idx = get<2, 1>(blk_coord); + int qo_len = get<0>(problem_shape); + int kv_len = get<1>(problem_shape); + int qo_segment_offset = get<0>(params_problem_shape).segment_offsets[batch_idx]; + int kv_segment_offset = get<1>(params_problem_shape).segment_offsets[batch_idx]; + + int mask_tile_count = Mask{}.get_trip_count(blk_coord, TileShape{}, problem_shape); + + using X = Underscore; + + // this one is only executed by one thread, no need to elect_one + + // Q1, K1, Q2, V1, K2, V2, K3, V3, ... + // two pipes: Q and KV + // from Memory (prod) to TensorCore (cons) + + Tensor mQ = params.tma_load_Q.get_tma_tensor(params.layout_Q.shape()); + Tensor mK = params.tma_load_K.get_tma_tensor(params.layout_K.shape()); + Tensor mV = params.tma_load_V.get_tma_tensor(params.layout_V.shape()); + + ThrMMA mma_qk = typename CollectiveMmaQK::TiledMma{}.get_slice(0); + ThrMMA mma_pv = typename CollectiveMmaPV::TiledMma{}.get_slice(0); + Tensor sQ = make_tensor(make_smem_ptr(storage.smem_q.data()), SmemLayoutQ{}); + Tensor sK = make_tensor(make_smem_ptr(storage.smem_k.data()), SmemLayoutK{}); + Tensor sV = make_tensor(make_smem_ptr(storage.smem_v.data()), SmemLayoutV{}); + + auto gQ = get_local_tile_tensor(mQ, select<0, 2>(TileShapeQK{}), qo_head_idx, qo_segment_offset, + qo_len); // (Q, D, _) + auto gK = get_local_tile_tensor(mK, select<1, 2>(TileShapeQK{}), qo_head_idx, kv_segment_offset, + kv_len); // (K, D, _) + auto gV = + get_local_tile_t_tensor(mV, select<1, 2>(TileShapePV{}), qo_head_idx, kv_segment_offset, + kv_len); // (K, D, _) + + int warp_idx = cutlass::canonical_warp_idx_sync(); + Tensor tSgQ_qdl = mma_qk.partition_A(gQ); + Tensor tSgK_kdl = mma_qk.partition_B(gK); + Tensor tOgV_dkl = mma_pv.partition_B(gV); + auto [tQgQ, tQsQ] = tma_partition(params.tma_load_Q, _0{}, Layout<_1>{}, group_modes<0, 3>(sQ), + group_modes<0, 3>(tSgQ_qdl)); // (TMA, q), (TMA, PIPE) + auto [tKgK, tKsK] = tma_partition(params.tma_load_K, _0{}, Layout<_1>{}, group_modes<0, 3>(sK), + group_modes<0, 3>(tSgK_kdl)); // (TMA, k), (TMA, PIPE) + auto [tVgV, tVsV] = tma_partition(params.tma_load_V, _0{}, Layout<_1>{}, group_modes<0, 3>(sV), + group_modes<0, 3>(tOgV_dkl)); // (TMA, k), (TMA, PIPE) + + // blk_coord in decomposed in terms of TileShape, not TileShapeQK + // As such, it needs to be transformed as + // (a,b,c): a -> 2*a (Q0) 2*a+1 (Q1) + // b -> 2*a (Ki i even) 2*a+1 (Ki i odd) + + uint32_t lane_predicate = cute::elect_one_sync(); + + // Q1 + int q0_index = 2 * get<0>(blk_coord); + int q1_index = 2 * get<0>(blk_coord) + 1; + pipeline_q.producer_acquire(pipeline_q_producer_state); + if (lane_predicate) { + auto tma_barrier = pipeline_q.producer_get_barrier(pipeline_q_producer_state); + copy(params.tma_load_Q.with(*tma_barrier, 0), tQgQ(_, q0_index), + tQsQ(_, pipeline_q_producer_state.index())); + } + ++pipeline_q_producer_state; + + // K1 + int k_index = 0; + pipeline_k.producer_acquire(pipeline_k_producer_state); + if (lane_predicate) { + auto tma_barrier = pipeline_k.producer_get_barrier(pipeline_k_producer_state); + copy(params.tma_load_K.with(*tma_barrier, 0), tKgK(_, k_index), + tKsK(_, pipeline_k_producer_state.index())); + } + ++pipeline_k_producer_state; + k_index += 1; + + // Q2 + pipeline_q.producer_acquire(pipeline_q_producer_state); + if (lane_predicate) { + auto tma_barrier = pipeline_q.producer_get_barrier(pipeline_q_producer_state); + copy(params.tma_load_Q.with(*tma_barrier, 0), tQgQ(_, q1_index), + tQsQ(_, pipeline_q_producer_state.index())); + } + ++pipeline_q_producer_state; + + // V1 + int v_index = 0; + pipeline_v.producer_acquire(pipeline_v_producer_state); + if (lane_predicate) { + auto tma_barrier = pipeline_v.producer_get_barrier(pipeline_v_producer_state); + copy(params.tma_load_V.with(*tma_barrier, 0), tVgV(_, v_index), + tVsV(_, pipeline_v_producer_state.index())); + } + ++pipeline_v_producer_state; + v_index += 1; + + // loop: + mask_tile_count -= 1; + for (; mask_tile_count > 0; mask_tile_count -= 1) { + // Ki + pipeline_k.producer_acquire(pipeline_k_producer_state); + if (lane_predicate) { + auto tma_barrier = pipeline_k.producer_get_barrier(pipeline_k_producer_state); + copy(params.tma_load_K.with(*tma_barrier, 0), tKgK(_, k_index), + tKsK(_, pipeline_k_producer_state.index())); + } + ++pipeline_k_producer_state; + k_index += 1; + + // Vi + pipeline_v.producer_acquire(pipeline_v_producer_state); + if (lane_predicate) { + auto tma_barrier = pipeline_v.producer_get_barrier(pipeline_v_producer_state); + copy(params.tma_load_V.with(*tma_barrier, 0), tVgV(_, v_index), + tVsV(_, pipeline_v_producer_state.index())); + } + ++pipeline_v_producer_state; + v_index += 1; + } + } +}; + +} // namespace cutlass::fmha::collective diff --git a/include/flashinfer/attention/blackwell/common/pow_2.hpp b/include/flashinfer/attention/blackwell/common/pow_2.hpp new file mode 100644 index 0000000000000000000000000000000000000000..64542402c8d8ca7f5c260c47a75813c799cfd6fd --- /dev/null +++ b/include/flashinfer/attention/blackwell/common/pow_2.hpp @@ -0,0 +1,89 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +#pragma once + +#include + +#include +#include + +namespace cutlass::fmha { + +struct Pow2 { + int n; + int log2_n; + + explicit CUTE_DEVICE Pow2(int n) : n(n) { +#ifdef __CUDA_ARCH__ + log2_n = __ffs(n) - 1; +#endif + } + + template + CUTE_HOST_DEVICE T operator*(T const& b) const { + return n * b; + } + + template + CUTE_HOST_DEVICE auto operator*(Int const&) const { + if constexpr (N & (N - 1) == 0) { + return Pow2{n * N}; + } + return n * N; + } +}; + +template +CUTE_HOST_DEVICE auto operator/(T const& a, Pow2 const& b) { + return a >> b.log2_n; +} + +template +CUTE_HOST_DEVICE auto operator%(T const& a, Pow2 const& b) { + return a & (b.n - 1); +} + +template +CUTE_HOST_DEVICE bool operator<(T const& a, Pow2 const& b) { + return a < b.n; +} + +CUTE_HOST_DEVICE void print(Pow2 const& a) { printf("2^%d", a.log2_n); } + +} // end namespace cutlass::fmha + +namespace cute { + +template <> +struct is_integral : true_type {}; + +} // end namespace cute diff --git a/include/flashinfer/attention/blackwell/device/fmha.hpp b/include/flashinfer/attention/blackwell/device/fmha.hpp new file mode 100644 index 0000000000000000000000000000000000000000..24d27298acfc8fdbab6cd10ef33f2afe1a8eb318 --- /dev/null +++ b/include/flashinfer/attention/blackwell/device/fmha.hpp @@ -0,0 +1,251 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! + \file + \brief An universal device layer for cutlass 3.x-style kernels. +*/ + +#pragma once + +// common +#include "cutlass/cutlass.h" +#include "cutlass/device_kernel.h" + +#if !defined(__CUDACC_RTC__) +#include "cutlass/cluster_launch.hpp" +#include "cutlass/trace.h" +#endif // !defined(__CUDACC_RTC__) + +//////////////////////////////////////////////////////////////////////////////// + +namespace cutlass::fmha::device { + +//////////////////////////////////////////////////////////////////////////////// +////////////////////////////// CUTLASS 3.x API ///////////////////////////////// +//////////////////////////////////////////////////////////////////////////////// + +template +class FMHA { + public: + using Kernel = Kernel_; + + static int const kThreadCount = Kernel::MaxThreadsPerBlock; + + /// Argument structure: User API + using Arguments = typename Kernel::Arguments; + /// Argument structure: Kernel API + using Params = typename Kernel::Params; + + private: + /// Kernel API parameters object + Params params_; + + bool is_initialized(bool set = false) { + static bool initialized = false; + if (set) initialized = true; + return initialized; + } + + public: + /// Access the Params structure + Params const& params() const { return params_; } + + /// Determines whether the GEMM can execute the given problem. + static Status can_implement(Arguments const& args) { + if (Kernel::can_implement(args)) { + return Status::kSuccess; + } else { + return Status::kInvalid; + } + } + + /// Gets the workspace size + static size_t get_workspace_size(Arguments const& args) { + size_t workspace_bytes = 0; + workspace_bytes += Kernel::get_workspace_size(args); + return workspace_bytes; + } + + /// Computes the grid shape + static dim3 get_grid_shape(Params const& params) { return Kernel::get_grid_shape(params); } + + /// Computes the maximum number of active blocks per multiprocessor + static int maximum_active_blocks(int /* smem_capacity */ = -1) { + CUTLASS_TRACE_HOST("FMHA::maximum_active_blocks()"); + int max_active_blocks = -1; + int smem_size = Kernel::SharedStorageSize; + + // first, account for dynamic smem capacity if needed + cudaError_t result; + if (smem_size >= (48 << 10)) { + CUTLASS_TRACE_HOST(" Setting smem size to " << smem_size); + result = cudaFuncSetAttribute(device_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size); + if (cudaSuccess != result) { + result = cudaGetLastError(); // to clear the error bit + CUTLASS_TRACE_HOST( + " cudaFuncSetAttribute() returned error: " << cudaGetErrorString(result)); + return -1; + } + } + + // query occupancy after setting smem size + result = cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &max_active_blocks, device_kernel, Kernel::MaxThreadsPerBlock, smem_size); + + if (cudaSuccess != result) { + result = cudaGetLastError(); // to clear the error bit + CUTLASS_TRACE_HOST(" cudaOccupancyMaxActiveBlocksPerMultiprocessor() returned error: " + << cudaGetErrorString(result)); + return -1; + } + + CUTLASS_TRACE_HOST(" max_active_blocks: " << max_active_blocks); + return max_active_blocks; + } + + /// Initializes GEMM state from arguments. + Status initialize(Arguments const& args, void* workspace = nullptr, + cudaStream_t stream = nullptr) { + CUTLASS_TRACE_HOST("FMHA::initialize() - workspace " + << workspace << ", stream: " << (stream ? "non-null" : "null")); + + // Initialize the workspace + Status status = Kernel::initialize_workspace(args, workspace, stream); + if (status != Status::kSuccess) { + return status; + } + + // Initialize the Params structure + params_ = Kernel::to_underlying_arguments(args, workspace); + + if (is_initialized()) return Status::kSuccess; + + // account for dynamic smem capacity if needed + int smem_size = Kernel::SharedStorageSize; + if (smem_size >= (48 << 10)) { + CUTLASS_TRACE_HOST(" Setting smem size to " << smem_size); + cudaError_t result = cudaFuncSetAttribute( + device_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size); + if (cudaSuccess != result) { + result = cudaGetLastError(); // to clear the error bit + CUTLASS_TRACE_HOST( + " cudaFuncSetAttribute() returned error: " << cudaGetErrorString(result)); + return Status::kErrorInternal; + } + } + + is_initialized(true); + + return Status::kSuccess; + } + + /// Update API is preserved in 3.0, but does not guarantee a lightweight update of params. + Status update(Arguments const& args, void* workspace = nullptr) { + CUTLASS_TRACE_HOST("FMHA()::update() - workspace: " << workspace); + + size_t workspace_bytes = get_workspace_size(args); + if (workspace_bytes > 0 && nullptr == workspace) { + return Status::kErrorWorkspaceNull; + } + + params_ = Kernel::to_underlying_arguments(args, workspace); + return Status::kSuccess; + } + + /// Primary run() entry point API that is static allowing users to create and manage their own + /// params. Supplied params struct must be construct by calling Kernel::to_underling_arguments() + static Status run(Params& params, cudaStream_t stream = nullptr, bool launch_with_pdl = true) { + CUTLASS_TRACE_HOST("FMHA::run()"); + dim3 const block = Kernel::get_block_shape(); + dim3 const grid = get_grid_shape(params); + + // configure smem size and carveout + int smem_size = Kernel::SharedStorageSize; + + Status launch_result; + // Use extended launch API only for mainloops that use it + if constexpr (Kernel::ArchTag::kMinComputeCapability >= 90) { + dim3 cluster(cute::size<0>(typename Kernel::ClusterShape{}), + cute::size<1>(typename Kernel::ClusterShape{}), + cute::size<2>(typename Kernel::ClusterShape{})); + void const* kernel = (void const*)device_kernel; + void* kernel_params[] = {¶ms}; + launch_result = ClusterLauncher::launch(grid, cluster, block, smem_size, stream, kernel, + kernel_params, launch_with_pdl); + } else { + launch_result = Status::kSuccess; + device_kernel<<>>(params); + } + + cudaError_t result = cudaGetLastError(); + if (cudaSuccess == result && Status::kSuccess == launch_result) { + return Status::kSuccess; + } else { + CUTLASS_TRACE_HOST(" Kernel launch failed. Reason: " << result); + return Status::kErrorInternal; + } + } + + // + // Non-static launch overloads that first create and set the internal params struct of this kernel + // handle. + // + + /// Launches the kernel after first constructing Params internal state from supplied arguments. + Status run(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr) { + Status status = initialize(args, workspace, stream); + if (Status::kSuccess == status) { + status = run(params_, stream); + } + return status; + } + + /// Launches the kernel after first constructing Params internal state from supplied arguments. + Status operator()(Arguments const& args, void* workspace = nullptr, + cudaStream_t stream = nullptr) { + return run(args, workspace, stream); + } + + /// Overload that allows a user to re-launch the same kernel without updating internal params + /// struct. + Status run(cudaStream_t stream = nullptr) { return run(params_, stream); } + + /// Overload that allows a user to re-launch the same kernel without updating internal params + /// struct. + Status operator()(cudaStream_t stream = nullptr) { return run(params_, stream); } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace cutlass::fmha::device + +//////////////////////////////////////////////////////////////////////////////// diff --git a/include/flashinfer/attention/blackwell/device/sm100_mla.hpp b/include/flashinfer/attention/blackwell/device/sm100_mla.hpp new file mode 100644 index 0000000000000000000000000000000000000000..9f724825693fdf6df3794e4abf12a610a2855dd0 --- /dev/null +++ b/include/flashinfer/attention/blackwell/device/sm100_mla.hpp @@ -0,0 +1,335 @@ +/*************************************************************************************************** + * Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! + \file + \brief An universal device layer for cutlass 3.x-style kernels. +*/ + +#pragma once + +// common +#include "cutlass/cutlass.h" +#include "cutlass/device_kernel.h" + +#if !defined(__CUDACC_RTC__) +#include "cutlass/cluster_launch.hpp" +#include "cutlass/trace.h" +#endif // !defined(__CUDACC_RTC__) + +#include "../kernel/sm100_fmha_mla_reduction.hpp" +#include "../kernel/sm100_fmha_mla_tma_warpspecialized.hpp" + +//////////////////////////////////////////////////////////////////////////////// + +namespace cutlass::fmha::device { + +using namespace cute; +using namespace cutlass::fmha::kernel; + +//////////////////////////////////////////////////////////////////////////////// +////////////////////////////// CUTLASS 3.x API ///////////////////////////////// +//////////////////////////////////////////////////////////////////////////////// + +template +class MLA { + public: + using Kernel = Kernel_; + + using ReductionKernel = cutlass::fmha::kernel::Sm100FmhaMlaReductionKernel< + typename Kernel::ElementOut, typename Kernel::ElementAcc, typename Kernel::ElementAcc, + Kernel::TileShapeH::value, Kernel::TileShapeL::value, 256 /*Max split*/ + >; + + /// Argument structure: User API + using KernelArguments = typename Kernel::Arguments; + using ReductionArguments = typename ReductionKernel::Arguments; + + using Arguments = KernelArguments; + + /// Argument structure: Kernel API + using KernelParams = typename Kernel::Params; + using ReductionParams = typename ReductionKernel::Params; + struct Params { + KernelParams fmha_params; + ReductionParams reduction_params; + }; + + private: + /// Kernel API parameters object + Params params_; + + bool is_initialized(bool set = false) { + static bool initialized = false; + if (set) initialized = true; + return initialized; + } + + static ReductionArguments to_reduction_args(Arguments const& args) { + auto [H, K, D, B] = args.problem_shape; + return ReductionArguments{nullptr, + args.epilogue.ptr_o, + nullptr, + args.epilogue.ptr_lse, + args.mainloop.softmax_scale, + B, + args.split_kv, + K, + args.mainloop.ptr_seq, + args.ptr_split_kv, + Kernel::TileShapeS::value}; + } + + public: + /// Access the Params structure + Params const& params() const { return params_; } + + static void set_split_kv(KernelArguments& args) { + if (args.split_kv >= 1) return; + auto [H, K, D, B] = args.problem_shape; + int sm_count = args.hw_info.sm_count; + int max_splits = ceil_div(K, 128); + int sms_per_batch = max(1, sm_count / B); + int split_heur = min(max_splits, sms_per_batch); + int waves = ceil_div(B * split_heur, sm_count); + int k_waves = ceil_div(max_splits, split_heur); + int split_wave_aware = ceil_div(max_splits, k_waves); + args.split_kv = split_wave_aware; + } + + /// Determines whether the GEMM can execute the given problem. + static Status can_implement(Arguments const& args) { + if (!Kernel::can_implement(args)) { + return Status::kInvalid; + } + if (!ReductionKernel::can_implement(to_reduction_args(args))) { + return Status::kInvalid; + } + return Status::kSuccess; + } + + /// Gets the workspace size + static size_t get_workspace_size(Arguments const& args) { + size_t workspace_bytes = 0; + workspace_bytes += Kernel::get_workspace_size(args); + workspace_bytes += ReductionKernel::get_workspace_size(to_reduction_args(args)); + return workspace_bytes; + } + + /// Computes the maximum number of active blocks per multiprocessor + static int maximum_active_blocks(int /* smem_capacity */ = -1) { + CUTLASS_TRACE_HOST("MLA::maximum_active_blocks()"); + int max_active_blocks = -1; + int smem_size = Kernel::SharedStorageSize; + + // first, account for dynamic smem capacity if needed + cudaError_t result; + if (smem_size >= (48 << 10)) { + CUTLASS_TRACE_HOST(" Setting smem size to " << smem_size); + result = cudaFuncSetAttribute(device_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size); + if (cudaSuccess != result) { + result = cudaGetLastError(); // to clear the error bit + CUTLASS_TRACE_HOST( + " cudaFuncSetAttribute() returned error: " << cudaGetErrorString(result)); + return -1; + } + } + + // query occupancy after setting smem size + result = cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &max_active_blocks, device_kernel, Kernel::MaxThreadsPerBlock, smem_size); + + if (cudaSuccess != result) { + result = cudaGetLastError(); // to clear the error bit + CUTLASS_TRACE_HOST(" cudaOccupancyMaxActiveBlocksPerMultiprocessor() returned error: " + << cudaGetErrorString(result)); + return -1; + } + + CUTLASS_TRACE_HOST(" max_active_blocks: " << max_active_blocks); + return max_active_blocks; + } + + /// Initializes GEMM state from arguments. + Status initialize(Arguments const& args, void* workspace = nullptr, + cudaStream_t stream = nullptr) { + CUTLASS_TRACE_HOST("MLA::initialize() - workspace " + << workspace << ", stream: " << (stream ? "non-null" : "null")); + + // Initialize the workspace + Status status = Kernel::initialize_workspace(args, workspace, stream); + if (status != Status::kSuccess) { + return status; + } + status = ReductionKernel::initialize_workspace(to_reduction_args(args), workspace, stream); + if (status != Status::kSuccess) { + return status; + } + KernelParams kernel_params = Kernel::to_underlying_arguments(args, workspace); + + ReductionArguments reduction_args = to_reduction_args(args); + if (reduction_args.split_kv > 1) { + reduction_args.ptr_oaccum = kernel_params.epilogue.ptr_o_acc; + reduction_args.ptr_lseaccum = kernel_params.epilogue.ptr_lse_acc; + } + ReductionParams reduction_params = + ReductionKernel::to_underlying_arguments(reduction_args, workspace); + // Initialize the Params structure + params_ = Params{kernel_params, reduction_params}; + + if (is_initialized()) return Status::kSuccess; + + // account for dynamic smem capacity if needed + // no dynamic smem is needed for reduction kernel + int smem_size = Kernel::SharedStorageSize; + if (smem_size >= (48 << 10)) { + CUTLASS_TRACE_HOST(" Setting smem size to " << smem_size); + cudaError_t result = cudaFuncSetAttribute( + device_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size); + if (cudaSuccess != result) { + result = cudaGetLastError(); // to clear the error bit + CUTLASS_TRACE_HOST( + " cudaFuncSetAttribute() returned error: " << cudaGetErrorString(result)); + return Status::kErrorInternal; + } + } + + is_initialized(true); + + return Status::kSuccess; + } + + /// Update API is preserved in 3.0, but does not guarantee a lightweight update of params. + Status update(Arguments const& args, void* workspace = nullptr) { + CUTLASS_TRACE_HOST("MLA()::update() - workspace: " << workspace); + + size_t workspace_bytes = get_workspace_size(args); + if (workspace_bytes > 0 && nullptr == workspace) { + return Status::kErrorWorkspaceNull; + } + + auto fmha_params = Kernel::to_underlying_arguments(args, workspace); + + ReductionArguments reduction_args = to_reduction_args(args); + if (reduction_args.split_kv > 1) { + reduction_args.ptr_oaccum = fmha_params.epilogue.ptr_o_acc; + reduction_args.ptr_lseaccum = fmha_params.epilogue.ptr_lse_acc; + } + ReductionParams reduction_params = + ReductionKernel::to_underlying_arguments(reduction_args, workspace); + // Initialize the Params structure + params_ = Params{fmha_params, reduction_params}; + + return Status::kSuccess; + } + + /// Primary run() entry point API that is static allowing users to create and manage their own + /// params. Supplied params struct must be construct by calling Kernel::to_underling_arguments() + static Status run(Params& params, cudaStream_t stream = nullptr) { + CUTLASS_TRACE_HOST("MLA::run()"); + dim3 const block = Kernel::get_block_shape(); + dim3 const grid = Kernel::get_grid_shape(params.fmha_params); + + // configure smem size and carveout + int smem_size = Kernel::SharedStorageSize; + + Status launch_result; + // Use extended launch API only for mainloops that use it + if constexpr (Kernel::ArchTag::kMinComputeCapability >= 90) { + dim3 cluster(cute::size<0>(typename Kernel::ClusterShape{}), + cute::size<1>(typename Kernel::ClusterShape{}), + cute::size<2>(typename Kernel::ClusterShape{})); + void const* kernel = (void const*)device_kernel; + void* kernel_params[] = {¶ms.fmha_params}; + launch_result = + ClusterLauncher::launch(grid, cluster, block, smem_size, stream, kernel, kernel_params); + } else { + launch_result = Status::kSuccess; + device_kernel<<>>(params.fmha_params); + } + + cudaError_t result = cudaGetLastError(); + if (cudaSuccess != result or Status::kSuccess != launch_result) { + // return Status::kSuccess; + CUTLASS_TRACE_HOST(" Kernel launch failed. Reason: " << result); + return Status::kErrorInternal; + } + if (params.reduction_params.split_kv > 1) { + // launch reduction kernel + dim3 const block = ReductionKernel::get_block_shape(); + dim3 const grid = ReductionKernel::get_grid_shape(params.reduction_params); + device_kernel<<>>(params.reduction_params); + cudaError_t result = cudaGetLastError(); + if (cudaSuccess == result) { + return Status::kSuccess; + } else { + CUTLASS_TRACE_HOST(" Kernel launch failed. Reason: " << result); + return Status::kErrorInternal; + } + } else { + return Status::kSuccess; + } + } + + // + // Non-static launch overloads that first create and set the internal params struct of this kernel + // handle. + // + + /// Launches the kernel after first constructing Params internal state from supplied arguments. + Status run(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr) { + Status status = initialize(args, workspace, stream); + if (Status::kSuccess == status) { + status = run(params_, stream); + } + return status; + } + + /// Launches the kernel after first constructing Params internal state from supplied arguments. + Status operator()(Arguments const& args, void* workspace = nullptr, + cudaStream_t stream = nullptr) { + return run(args, workspace, stream); + } + + /// Overload that allows a user to re-launch the same kernel without updating internal params + /// struct. + Status run(cudaStream_t stream = nullptr) { return run(params_, stream); } + + /// Overload that allows a user to re-launch the same kernel without updating internal params + /// struct. + Status operator()(cudaStream_t stream = nullptr) { return run(params_, stream); } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace cutlass::fmha::device + +//////////////////////////////////////////////////////////////////////////////// diff --git a/include/flashinfer/attention/blackwell/fmha_cutlass_sm100.cuh b/include/flashinfer/attention/blackwell/fmha_cutlass_sm100.cuh new file mode 100644 index 0000000000000000000000000000000000000000..66d8b4fb913d7978a68244cfc737e11ab3680c20 --- /dev/null +++ b/include/flashinfer/attention/blackwell/fmha_cutlass_sm100.cuh @@ -0,0 +1,178 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include + +#include "../../allocator.h" +#include "collective/fmha_fusion.hpp" +#include "collective/sm100_fmha_fwd_epilogue_tma_warpspecialized.hpp" +#include "collective/sm100_fmha_fwd_mainloop_tma_warpspecialized.hpp" +#include "cute/tensor.hpp" +#include "cutlass/cutlass.h" +#include "cutlass/kernel_hardware_info.h" +#include "cutlass/util/command_line.h" +#include "cutlass/util/device_memory.h" +#include "cutlass/util/distribution.h" +#include "cutlass/util/reference/device/tensor_fill.h" +#include "device/fmha.hpp" +#include "kernel/fmha_tile_scheduler.hpp" +#include "kernel/sm100_fmha_fwd_kernel_tma_warpspecialized.hpp" + +namespace flashinfer { + +using namespace cute; +using namespace cutlass::fmha::collective; +using namespace cutlass::fmha::kernel; +using namespace cutlass::fmha::device; + +template +struct FwdRunner { + using Element = DTypeIn; + using ElementAccumulatorQK = float; + using ElementAccumulatorPV = float; + using ElementOut = DTypeOut; + + // Q K D ((H_R, H_KV), B) + using ProblemShapeVarlen = + cute::tuple, int>>; + + using StrideQ = cute::tuple>; // Q D (H_G H_R) + using StrideK = cute::tuple>; // K D (H_G H_R) + using StrideV = cute::tuple<_1, int, cute::tuple<_0, int>>; // D V (H_G H_R) + // NOTE(Zihao): use markus's trick for tma store + using StrideO = + cute::tuple, int>>; // Q D (H_G H_R) CUMULATIVE_Q + using StrideLSE = cute::tuple>; // Q (H_G H_R) + + using Mainloop = cutlass::fmha::collective::Sm100FmhaFwdMainloopTmaWarpspecialized< + Element, ElementAccumulatorQK, ElementAccumulatorPV, TileShapeQK, TileShapePV, StrideQ, + StrideK, StrideV, ActiveMask>; + using Epilogue = cutlass::fmha::collective::Sm100FmhaFwdEpilogueTmaWarpspecialized< + ElementOut, ElementAccumulatorPV, typename Mainloop::TileShapePV>; + using Operation = + cutlass::fmha::device::FMHA>; + using LayoutQ = typename Mainloop::LayoutQ; + using LayoutK = typename Mainloop::LayoutK; + using LayoutV = typename Mainloop::LayoutV; + using LayoutO = typename Epilogue::LayoutO; + using LayoutLSE = typename Epilogue::LayoutLSE; + + static cudaError_t run(void* workspace_buffer, DTypeIn* q, DTypeIn* k, DTypeIn* v, + IdType* qo_segment_offsets, IdType* kv_segment_offsets, + IdType* work_indptr, IdType* qo_tile_indices, IdType* qo_head_indices, + IdType* batch_indices, DTypeOut* o, float* maybe_lse, int mask_mode_code, + double sm_scale, int num_qo_heads, int num_kv_heads, int head_dim_qk, + int head_dim_vo, int q_stride_n, int q_stride_h, int k_stride_n, + int k_stride_h, int v_stride_n, int v_stride_h, int batch_size, + int total_qo_len, int total_kv_len, int max_qo_len, cudaStream_t stream) { + cutlass::KernelHardwareInfo hw_info; + hw_info.device_id = 0; + hw_info.sm_count = + cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id); + + StrideQ stride_Q; + StrideK stride_K; + StrideV stride_V; + StrideO stride_O; + StrideLSE stride_LSE; + + int h_r = num_qo_heads / num_kv_heads; + assert(num_qo_heads % num_kv_heads == 0); + ProblemShapeVarlen problem_shape = cute::make_tuple( + VariableLength{qo_segment_offsets}, VariableLength{kv_segment_offsets}, head_dim_qk, + cute::make_tuple(cute::make_tuple(h_r, num_kv_heads), batch_size)); + + stride_Q = make_stride(q_stride_n, _1{}, make_stride(q_stride_h, h_r * q_stride_h)); + stride_O = make_stride( + num_qo_heads * head_dim_vo, _1{}, + make_stride(make_stride(head_dim_vo, h_r * head_dim_vo), num_qo_heads * head_dim_vo)); + stride_K = make_stride(k_stride_n, _1{}, make_stride(_0{}, k_stride_h)); + stride_V = make_stride(_1{}, v_stride_n, make_stride(_0{}, v_stride_h)); + stride_LSE = make_stride(num_qo_heads, make_stride(_1{}, h_r)); + + auto shape_Q = make_shape(total_qo_len, head_dim_qk, make_shape(h_r, num_kv_heads)); + auto shape_O = make_shape(max_qo_len, head_dim_vo, + make_shape(make_shape(h_r, num_kv_heads), max_qo_len + total_qo_len)); + auto shape_K = make_shape(total_kv_len, head_dim_qk, make_shape(h_r, num_kv_heads)); + auto shape_V = make_shape(head_dim_vo, total_kv_len, make_shape(h_r, num_kv_heads)); + auto shape_LSE = make_shape(total_qo_len, make_shape(h_r, num_kv_heads)); + + LayoutQ layout_Q = make_layout(shape_Q, stride_Q); + LayoutK layout_K = make_layout(shape_K, stride_K); + LayoutV layout_V = make_layout(shape_V, stride_V); + LayoutO layout_O = make_layout(shape_O, stride_O); + LayoutLSE layout_LSE = make_layout(shape_LSE, stride_LSE); + + typename Operation::Arguments arguments{ + problem_shape, + {q, layout_Q, k, layout_K, v, layout_V, sm_scale}, + {o - max_qo_len * get<0>(stride_O), layout_O, maybe_lse, layout_LSE, max_qo_len}, + {work_indptr, qo_tile_indices, qo_head_indices, batch_indices}, + hw_info}; + + Operation op; + + // NOTE(Zihao): workspace size is not used at this moment + size_t workspace_size = 0; + workspace_size = Operation::get_workspace_size(arguments); + AlignedAllocator allocator(workspace_buffer, workspace_size); + uint8_t* workspace_ptr = + allocator.aligned_alloc(workspace_size, 16, "fmha_cutlass_sm100_workspace"); + + cutlass::Status status = cutlass::Status::kSuccess; + status = op.can_implement(arguments); + if (status != cutlass::Status::kSuccess) { + std::cerr << "This kernel is not supported. Last CUDA error is: " + << cudaGetErrorString(cudaGetLastError()) << std::endl; + } + + status = op.initialize(arguments, workspace_ptr); + if (status != cutlass::Status::kSuccess) { + std::cerr << "Failed to initialize the CUTLASS kernel. Last CUDA error is: " + << cudaGetErrorString(cudaGetLastError()) << std::endl; + } + + // Run + status = op.run(stream); + if (status != cutlass::Status::kSuccess) { + std::cerr << "Failed to launch the CUTLASS kernel. Last CUDA error is: " + << cudaGetErrorString(cudaGetLastError()) << std::endl; + } + return cudaSuccess; + } +}; + +template +cudaError_t run_fmha_fwd(void* workspace_buffer, DTypeIn* q, DTypeIn* k, DTypeIn* v, + IdType* qo_segment_offsets, IdType* kv_segment_offsets, + IdType* work_indptr, IdType* qo_tile_indices, IdType* qo_head_indices, + IdType* batch_indices, DTypeOut* o, float* maybe_lse, int mask_mode_code, + double sm_scale, int num_qo_heads, int num_kv_heads, int head_dim_qk, + int head_dim_vo, int q_stride_n, int q_stride_h, int k_stride_n, + int k_stride_h, int v_stride_n, int v_stride_h, int batch_size, + int total_qo_len, int total_kv_len, int max_qo_len, cudaStream_t stream) { + return FwdRunner::run( + workspace_buffer, q, k, v, qo_segment_offsets, kv_segment_offsets, work_indptr, + qo_tile_indices, qo_head_indices, batch_indices, o, maybe_lse, mask_mode_code, sm_scale, + num_qo_heads, num_kv_heads, head_dim_qk, head_dim_vo, q_stride_n, q_stride_h, k_stride_n, + k_stride_h, v_stride_n, v_stride_h, batch_size, total_qo_len, total_kv_len, max_qo_len, + stream); +} + +}; // namespace flashinfer diff --git a/include/flashinfer/attention/blackwell/kernel/fmha_options.hpp b/include/flashinfer/attention/blackwell/kernel/fmha_options.hpp new file mode 100644 index 0000000000000000000000000000000000000000..5b32c079cf57c14d7dea6656573d06813f9bdd5d --- /dev/null +++ b/include/flashinfer/attention/blackwell/kernel/fmha_options.hpp @@ -0,0 +1,78 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +#pragma once + +#include "cutlass/cutlass.h" + +namespace cutlass::fmha::kernel { + +template +struct find_option; + +template +struct find_option { + using option_value = Default; +}; + +template +struct find_option + : std::conditional_t > {}; + +template +using find_option_t = typename find_option::option_value; + +enum class Tag { + kIsPersistent, + kNumMmaWarpGroups, + kLoadsQSeparately, + + kIsMainloopLocked, + kIsEpilogueLocked, + + kStagesQ, + kStagesKV, + + kEpilogueKind, + + kBlocksPerSM, + kClusterM, + + kAccQK +}; + +template +struct Option { + static constexpr auto tag = kTag; + using option_value = Value; +}; + +} // namespace cutlass::fmha::kernel diff --git a/include/flashinfer/attention/blackwell/kernel/fmha_tile_scheduler.hpp b/include/flashinfer/attention/blackwell/kernel/fmha_tile_scheduler.hpp new file mode 100644 index 0000000000000000000000000000000000000000..d1369b433f9632fe63f882d30f9eabc046a6914c --- /dev/null +++ b/include/flashinfer/attention/blackwell/kernel/fmha_tile_scheduler.hpp @@ -0,0 +1,112 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +#pragma once + +#include "cutlass/cutlass.h" +#include "cutlass/fast_math.h" +#include "cutlass/kernel_hardware_info.h" + +namespace cutlass::fmha::kernel { + +struct HostPrecomputedTileScheduler { + struct Arguments { + int* work_indptr; + int* qo_tile_indices; + int* qo_head_indices; + int* batch_indices; + }; + + struct Params { + int* work_indptr; + int* qo_tile_indices; + int* qo_head_indices; + int* batch_indices; + int num_sm; + }; + + Params params; + int work_ptr; + int work_ptr_end; + int qo_tile_idx; + int batch_idx; + int qo_head_idx; + bool is_valid_; + + CUTLASS_DEVICE + HostPrecomputedTileScheduler(Params const& params) { + this->params = params; + work_ptr = params.work_indptr[blockIdx.x]; + work_ptr_end = params.work_indptr[blockIdx.x + 1]; + if (work_ptr < work_ptr_end) { + qo_tile_idx = params.qo_tile_indices[work_ptr]; + batch_idx = params.batch_indices[work_ptr]; + qo_head_idx = params.qo_head_indices[work_ptr]; + } else { + qo_tile_idx = 0; + batch_idx = 0; + qo_head_idx = 0; + } + is_valid_ = true; + } + + static Params to_underlying_arguments(Arguments const& args, KernelHardwareInfo hw_info) { + return {args.work_indptr, args.qo_tile_indices, args.qo_head_indices, args.batch_indices, + hw_info.sm_count}; + } + + static dim3 get_grid_shape(Params const& params) { + dim3 grid(params.num_sm); + return grid; + } + + CUTLASS_DEVICE + bool is_valid() const { return is_valid_; } + + CUTLASS_DEVICE + auto get_block_coord() { + return make_coord(qo_tile_idx, _0{}, make_coord(qo_head_idx, batch_idx)); + } + + CUTLASS_DEVICE + HostPrecomputedTileScheduler& operator++() { + work_ptr++; + is_valid_ = work_ptr < work_ptr_end; + if (is_valid_) { + qo_tile_idx = params.qo_tile_indices[work_ptr]; + batch_idx = params.batch_indices[work_ptr]; + qo_head_idx = params.qo_head_indices[work_ptr]; + } + return *this; + } +}; + +} // namespace cutlass::fmha::kernel diff --git a/include/flashinfer/attention/blackwell/kernel/gather_tensor.hpp b/include/flashinfer/attention/blackwell/kernel/gather_tensor.hpp new file mode 100644 index 0000000000000000000000000000000000000000..17e4475f7b46b09d9eb285ad7b699d6ba3c5ec4d --- /dev/null +++ b/include/flashinfer/attention/blackwell/kernel/gather_tensor.hpp @@ -0,0 +1,190 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include "cute/layout.hpp" +#include "cute/tensor.hpp" +#include "cute/util/print.hpp" + +namespace example { + +using namespace cute; + +// Empty type used to disable gather/scatter for a GEMM argument +struct NoGather { + template + NoGather(Ts...){}; +}; + +/// Function object that applies an index to its argument +template +struct IndexedGather { + CUTE_HOST_DEVICE constexpr IndexedGather(Index const* indices = {}) : indices_(indices) {} + + template + CUTE_HOST_DEVICE constexpr Index operator()(I i) const { + return indices_[i]; + } + + CUTE_HOST_DEVICE friend void print(IndexedGather const& s) { cute::print("Indexed"); } + + Index const* indices_; +}; + +/// Function object that applies a stride to its argument +/// Example: StridedFunc gathers every other row/column +template +struct StridedGather { + CUTE_HOST_DEVICE constexpr StridedGather(Stride stride = {}) : stride_(stride) {} + + template + CUTE_HOST_DEVICE constexpr auto operator()(I i) const { + return i * stride_; + } + + CUTE_HOST_DEVICE friend void print(StridedGather const& s) { + cute::print("Strided{"); + print(s.stride_); + cute::print("}"); + } + + Stride stride_; +}; + +/// Custom stride object that applies a function followed by a stride +template +struct CustomStride { + CUTE_HOST_DEVICE constexpr CustomStride(Func const& func, Stride const& stride) + : func_(func), stride_(stride) {} + + template + CUTE_HOST_DEVICE constexpr friend auto operator*(I i, CustomStride const& s) { + return s.func_(i) * s.stride_; + } + + template + CUTE_HOST_DEVICE constexpr friend auto operator*(CustomStride const& s, I i) { + return s.func_(i) * s.stride_; + } + + CUTE_HOST_DEVICE friend void print(CustomStride const& s) { + cute::print("Custom{"); + print(s.func_); + cute::print(","); + print(s.stride_); + cute::print("}"); + } + + template + CUTE_HOST_DEVICE constexpr friend auto safe_div(CustomStride const& s, Div const& div) { + return CustomStride(s.func_, + safe_div(s.stride_, div)); + } + + // Circumvent the requirement on make_layout that shape and stride are integral + template + CUTE_HOST_DEVICE constexpr friend auto make_layout(Shape const& shape, + CustomStride const& stride) { + return Layout(shape, stride); + } + + Func func_; + Stride stride_; +}; + +template +CUTLASS_HOST_DEVICE auto make_custom_stride_layout(Stride const& stride, Func&& func) { + // Use a dummy shape and replace the first non-unit stride with a custom gather stride + auto idx = find_if(stride, [](auto x) { return not is_constant<1, decltype(x)>{}; }); + constexpr int I = decltype(idx)::value; + return make_layout(repeat_like(stride, _1{}), + replace(stride, CustomStride{static_cast(func), get(stride)})); +} + +/// Helper function to optionally create a gather tensor +template +CUTLASS_HOST_DEVICE auto make_gather_tensor(Iterator iter, Shape const& shape, Stride const& stride, + Func&& func) { + if constexpr (not cutlass::platform::is_same, NoGather>::value) { + Layout matrix_layout = make_identity_layout(shape); + auto offset = as_arithmetic_tuple(repeat_like(shape, _0{})); + Layout gather_layout = make_custom_stride_layout(stride, static_cast(func)); + return make_tensor(iter, ComposedLayout{gather_layout, offset, matrix_layout}); + } else { + return make_tensor(iter, shape, stride); + } +} + +} // namespace example + +namespace cute { + +template +CUTE_HOST_DEVICE constexpr auto upcast(Shape const& shape, Stride const& stride) { + if constexpr (is_tuple::value) { + return transform_layout(shape, stride, + [](auto const& s, auto const& d) { return upcast(s, d); }); + } else if constexpr (is_scaled_basis::value) { + if constexpr (Stride::mode() == I) { + return make_layout(ceil_div(shape, Int{}), ceil_div(stride, Int{})); + } else { + return make_layout(shape, stride); + } + } else { + return upcast(shape, stride); + } + + CUTE_GCC_UNREACHABLE; +} + +template +CUTE_HOST_DEVICE constexpr auto upcast( + ComposedLayout, Offset, Layout> const& layout) { + // Find index of the stride-1 mode - that is the only one that requires updating inner shape and + // offset + auto idx = + find_if(layout.layout_a().stride(), [](auto x) { return is_constant<1, decltype(x)>{}; }); + constexpr int I = decltype(idx)::value; + + // Upcast the outer layout (works as expected) + auto outer = upcast(layout.layout_a()); + + // Upcast the accumulated offset along stride-1 mode + auto offset = + as_arithmetic_tuple(replace(layout.offset(), upcast(get(layout.offset())))); + + // Upcast the inner layout's shape along stride-1 mode + auto inner = upcast(layout.layout_b().shape(), layout.layout_b().stride()); + + return composition(outer, offset, inner); +} + +} // namespace cute diff --git a/include/flashinfer/attention/blackwell/kernel/sm100_fmha_fwd_kernel_tma_warpspecialized.hpp b/include/flashinfer/attention/blackwell/kernel/sm100_fmha_fwd_kernel_tma_warpspecialized.hpp new file mode 100644 index 0000000000000000000000000000000000000000..d6e913a319aec5fb0eb87491a60b8c410b39a318 --- /dev/null +++ b/include/flashinfer/attention/blackwell/kernel/sm100_fmha_fwd_kernel_tma_warpspecialized.hpp @@ -0,0 +1,532 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +#include "../collective/fmha_common.hpp" +#include "../collective/fmha_fusion.hpp" +#include "cute/arch/tmem_allocator_sm100.hpp" +#include "cute/layout.hpp" +#include "cutlass/arch/arch.h" +#include "cutlass/cutlass.h" +#include "cutlass/kernel_hardware_info.h" +#include "cutlass/pipeline/pipeline.hpp" +#include "fmha_options.hpp" +#include "fmha_tile_scheduler.hpp" + +namespace cutlass::fmha::kernel { + +using namespace cute; +using namespace cutlass::fmha::collective; + +struct Sm100FmhaCtxKernelWarpspecializedSchedule { + enum class WarpRole { Softmax0, Softmax1, Correction, MMA, Load, Epilogue, Empty }; + + static constexpr WarpRole warp_idx_to_WarpRole(int warp_idx) { + int wg_idx = warp_idx / 4; // warp_idx + if (wg_idx == 0) return WarpRole::Softmax0; // 0 - 3 + if (wg_idx == 1) return WarpRole::Softmax1; // 4 - 7 + if (wg_idx == 2) return WarpRole::Correction; // 8 - 11 + if (warp_idx == 12) return WarpRole::MMA; // 12 + if (warp_idx == 13) return WarpRole::Load; // 13 + if (warp_idx == 14) return WarpRole::Epilogue; // 14 + return WarpRole::Empty; // 15 + } + + static const int NumWarpsSoftmax = 4; + static const int NumWarpsCorrection = 4; + static const int NumWarpsEpilogue = 1; + static const int NumWarpsLoad = 1; + + static const bool kDebugUsingPrintf = false; + static const int NumRegsSoftmax = 192; + static const int NumRegsCorrection = 96 - (kDebugUsingPrintf ? 16 : 0); + static const int NumRegsOther = 32 + (kDebugUsingPrintf ? 16 : 0); + static const int NumRegsEmpty = 24; + + static const int NumWarps = 16; +}; + +template +struct Sm100FmhaFwdKernelTmaWarpspecialized { + using TileShape = typename CollectiveMainloop::TileShape; + using ProblemShape = ProblemShapeIn; + + using WarpRole = typename KernelSchedule::WarpRole; + + constexpr WarpRole warp_idx_to_WarpRole(int warp_idx) { + return KernelSchedule::warp_idx_to_WarpRole(warp_idx); + } + + static const int NumWarpsSoftmax = KernelSchedule::NumWarpsSoftmax; + static const int NumWarpsCorrection = KernelSchedule::NumWarpsCorrection; + static const int NumWarpsEpilogue = KernelSchedule::NumWarpsEpilogue; + static const int NumWarpsLoad = KernelSchedule::NumWarpsLoad; + + static const int NumRegsSoftmax = KernelSchedule::NumRegsSoftmax; + static const int NumRegsCorrection = KernelSchedule::NumRegsCorrection; + static const int NumRegsOther = KernelSchedule::NumRegsOther; + static const int NumRegsEmpty = 24; + + static const int NumWarps = KernelSchedule::NumWarps; + + using ClusterShape = typename CollectiveMainloop::ClusterShape; + + using TmemAllocator = cute::TMEM::Allocator1Sm; + + struct SharedStorage { + union { + typename CollectiveMainloop::TensorStorage mainloop; + typename CollectiveEpilogue::TensorStorage epilogue; + }; + + struct PipelineStorage { + alignas(16) typename CollectiveMainloop::PipelineQ::SharedStorage load_q; + alignas(16) typename CollectiveMainloop::PipelineK::SharedStorage load_k; + alignas(16) typename CollectiveMainloop::PipelineV::SharedStorage load_v; + alignas(16) typename CollectiveMainloop::PipelineS::SharedStorage mma_s0; + alignas(16) typename CollectiveMainloop::PipelineS::SharedStorage mma_s1; + alignas(16) typename CollectiveMainloop::PipelineC::SharedStorage s0_corr; + alignas(16) typename CollectiveMainloop::PipelineC::SharedStorage s1_corr; + alignas(16) typename CollectiveMainloop::PipelineO::SharedStorage mma_corr; + alignas(16) typename CollectiveMainloop::PipelineE::SharedStorage corr_epi; + alignas(16) typename CollectiveMainloop::OrderBarrierSoftmax::SharedStorage order_s01; + } pipelines; + + cutlass::arch::ClusterBarrier barrier_O; + uint32_t tmem_base_ptr; + }; + + static constexpr int SharedStorageSize = sizeof(SharedStorage); + + struct Arguments { + ProblemShape problem_shape; + typename CollectiveMainloop::Arguments mainloop; + typename CollectiveEpilogue::Arguments epilogue; + typename TileScheduler::Arguments tile_scheduler; + cutlass::KernelHardwareInfo hw_info; + }; + + struct Params { + ProblemShape problem_shape; + typename CollectiveMainloop::Params mainloop; + typename CollectiveEpilogue::Params epilogue; + typename TileScheduler::Params tile_scheduler; + }; + + static const int MinBlocksPerMultiprocessor = 1; + static const int MaxThreadsPerBlock = NumWarps * cutlass::NumThreadsPerWarp; + using ArchTag = cutlass::arch::Sm100; + + static size_t get_workspace_size(Arguments const& args) { return 0; } + static cutlass::Status initialize_workspace(Arguments const&, void*, cudaStream_t) { + return cutlass::Status::kSuccess; + } + + static bool can_implement(Arguments const& args) { + return CollectiveMainloop::can_implement(args.problem_shape, args.mainloop); + } + + static dim3 get_grid_shape(Params const& params) { + return TileScheduler::get_grid_shape(params.tile_scheduler); + } + + static dim3 get_block_shape() { + dim3 block(MaxThreadsPerBlock, 1, 1); + return block; + } + + static Params to_underlying_arguments(Arguments const& args, void* workspace) { + return Params{ + args.problem_shape, + CollectiveMainloop::to_underlying_arguments(args.problem_shape, args.mainloop, workspace), + CollectiveEpilogue::to_underlying_arguments(args.problem_shape, args.epilogue, workspace), + TileScheduler::to_underlying_arguments(args.tile_scheduler, args.hw_info)}; + } + + CUTLASS_DEVICE auto apply_batch(const Params& params, ProblemShape const& problem_shape, + int batch_idx) { + return apply_variable_length(params.problem_shape, batch_idx); + } + + CUTLASS_DEVICE void operator()(const Params& params, char* smem) { +#if (__CUDACC_VER_MAJOR__ >= 12 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.wait;"); +#endif + + TileScheduler tile_scheduler{params.tile_scheduler}; + + int warp_idx = cutlass::canonical_warp_idx_sync(); + auto role = warp_idx_to_WarpRole(warp_idx); + uint32_t lane_predicate = cute::elect_one_sync(); + + if (role == WarpRole::Load && lane_predicate) { + CollectiveMainloop::prefetch_tma_descriptors(params.mainloop); + } + + if (role == WarpRole::Epilogue && lane_predicate) { + CollectiveEpilogue::prefetch_tma_descriptors(params.epilogue); + } + + SharedStorage& shared_storage = *reinterpret_cast(smem); + + typename CollectiveMainloop::PipelineQ::Params pipeline_load_q_params; + if (role == WarpRole::Load) { + pipeline_load_q_params.role = CollectiveMainloop::PipelineQ::ThreadCategory::Producer; + } + if (role == WarpRole::MMA) { + pipeline_load_q_params.role = CollectiveMainloop::PipelineQ::ThreadCategory::Consumer; + } + pipeline_load_q_params.is_leader = lane_predicate && (role == WarpRole::Load); + pipeline_load_q_params.transaction_bytes = CollectiveMainloop::TransactionBytesLoadQ; + typename CollectiveMainloop::PipelineQ pipeline_load_q( + shared_storage.pipelines.load_q, pipeline_load_q_params, ClusterShape{}, cute::true_type{}, + /*mask calc*/ cute::false_type{}); + + typename CollectiveMainloop::PipelineK::Params pipeline_load_k_params; + typename CollectiveMainloop::PipelineV::Params pipeline_load_v_params; + if (role == WarpRole::Load) { + pipeline_load_k_params.role = CollectiveMainloop::PipelineK::ThreadCategory::Producer; + pipeline_load_v_params.role = CollectiveMainloop::PipelineV::ThreadCategory::Producer; + } + if (role == WarpRole::MMA) { + pipeline_load_k_params.role = CollectiveMainloop::PipelineK::ThreadCategory::Consumer; + pipeline_load_v_params.role = CollectiveMainloop::PipelineV::ThreadCategory::Consumer; + } + pipeline_load_k_params.is_leader = lane_predicate && (role == WarpRole::Load); + pipeline_load_v_params.is_leader = lane_predicate && (role == WarpRole::Load); + pipeline_load_k_params.transaction_bytes = CollectiveMainloop::TransactionBytesLoadK; + pipeline_load_v_params.transaction_bytes = CollectiveMainloop::TransactionBytesLoadV; + typename CollectiveMainloop::PipelineK pipeline_load_k( + shared_storage.pipelines.load_k, pipeline_load_k_params, ClusterShape{}, + /*barrier init*/ cute::true_type{}, /*mask calc*/ cute::false_type{}); + typename CollectiveMainloop::PipelineV pipeline_load_v( + shared_storage.pipelines.load_v, pipeline_load_v_params, ClusterShape{}, + /*barrier init*/ cute::true_type{}, /*mask calc*/ cute::false_type{}); + + typename CollectiveMainloop::PipelineS::Params pipeline_mma_s0_params; + if (role == WarpRole::MMA) { + pipeline_mma_s0_params.role = CollectiveMainloop::PipelineS::ThreadCategory::Producer; + } + if (role == WarpRole::Softmax0) { + pipeline_mma_s0_params.role = CollectiveMainloop::PipelineS::ThreadCategory::Consumer; + } + pipeline_mma_s0_params.consumer_arv_count = NumWarpsSoftmax * cutlass::NumThreadsPerWarp; + typename CollectiveMainloop::PipelineS pipeline_mma_s0( + shared_storage.pipelines.mma_s0, pipeline_mma_s0_params, ClusterShape{}, + /*barrier init*/ cute::true_type{}, /*mask calc*/ cute::false_type{}); + + typename CollectiveMainloop::PipelineS::Params pipeline_mma_s1_params; + if (role == WarpRole::MMA) { + pipeline_mma_s1_params.role = CollectiveMainloop::PipelineS::ThreadCategory::Producer; + } + if (role == WarpRole::Softmax1) { + pipeline_mma_s1_params.role = CollectiveMainloop::PipelineS::ThreadCategory::Consumer; + } + pipeline_mma_s1_params.consumer_arv_count = NumWarpsSoftmax * cutlass::NumThreadsPerWarp; + typename CollectiveMainloop::PipelineS pipeline_mma_s1( + shared_storage.pipelines.mma_s1, pipeline_mma_s1_params, ClusterShape{}, + /*barrier init*/ cute::true_type{}, /*mask calc*/ cute::false_type{}); + + typename CollectiveMainloop::PipelineC::Params pipeline_s0_corr_params; + if (role == WarpRole::Softmax0) { + pipeline_s0_corr_params.role = CollectiveMainloop::PipelineC::ThreadCategory::Producer; + } + if (role == WarpRole::Correction) { + pipeline_s0_corr_params.role = CollectiveMainloop::PipelineC::ThreadCategory::Consumer; + } + pipeline_s0_corr_params.producer_arv_count = NumWarpsSoftmax * cutlass::NumThreadsPerWarp; + pipeline_s0_corr_params.consumer_arv_count = NumWarpsCorrection * cutlass::NumThreadsPerWarp; + typename CollectiveMainloop::PipelineC pipeline_s0_corr(shared_storage.pipelines.s0_corr, + pipeline_s0_corr_params, + /*barrier init*/ cute::true_type{}); + + typename CollectiveMainloop::PipelineC::Params pipeline_s1_corr_params; + if (role == WarpRole::Softmax1) { + pipeline_s1_corr_params.role = CollectiveMainloop::PipelineC::ThreadCategory::Producer; + } + if (role == WarpRole::Correction) { + pipeline_s1_corr_params.role = CollectiveMainloop::PipelineC::ThreadCategory::Consumer; + } + pipeline_s1_corr_params.producer_arv_count = NumWarpsSoftmax * cutlass::NumThreadsPerWarp; + pipeline_s1_corr_params.consumer_arv_count = NumWarpsCorrection * cutlass::NumThreadsPerWarp; + typename CollectiveMainloop::PipelineC pipeline_s1_corr(shared_storage.pipelines.s1_corr, + pipeline_s1_corr_params, + /*barrier init*/ cute::true_type{}); + + typename CollectiveMainloop::PipelineO::Params pipeline_mma_corr_params; + if (role == WarpRole::MMA) { + pipeline_mma_corr_params.role = CollectiveMainloop::PipelineO::ThreadCategory::Producer; + } + if (role == WarpRole::Correction) { + pipeline_mma_corr_params.role = CollectiveMainloop::PipelineO::ThreadCategory::Consumer; + } + pipeline_mma_corr_params.consumer_arv_count = NumWarpsCorrection * cutlass::NumThreadsPerWarp; + typename CollectiveMainloop::PipelineO pipeline_mma_corr( + shared_storage.pipelines.mma_corr, pipeline_mma_corr_params, ClusterShape{}, + /*barrier init*/ cute::true_type{}, /*mask calc*/ cute::false_type{}); + + typename CollectiveMainloop::PipelineE::Params pipeline_corr_epi_params; + if (role == WarpRole::Correction) { + pipeline_corr_epi_params.role = CollectiveMainloop::PipelineE::ThreadCategory::Producer; + } + if (role == WarpRole::Epilogue) { + pipeline_corr_epi_params.role = CollectiveMainloop::PipelineE::ThreadCategory::Consumer; + } + pipeline_corr_epi_params.producer_arv_count = NumWarpsCorrection * cutlass::NumThreadsPerWarp; + pipeline_corr_epi_params.consumer_arv_count = NumWarpsEpilogue * cutlass::NumThreadsPerWarp; + typename CollectiveMainloop::PipelineE pipeline_corr_epi(shared_storage.pipelines.corr_epi, + pipeline_corr_epi_params, + /*barrier init*/ cute::true_type{}); + + typename CollectiveMainloop::OrderBarrierSoftmax::Params params_order_s01; + params_order_s01.group_id = role == WarpRole::Softmax1 ? 1 : 0; + params_order_s01.group_size = NumWarpsSoftmax * cutlass::NumThreadsPerWarp; + typename CollectiveMainloop::OrderBarrierSoftmax order_s01(shared_storage.pipelines.order_s01, + params_order_s01); + + TmemAllocator tmem_allocator; + + if (role == WarpRole::Load && lane_predicate) { + shared_storage.barrier_O.init(/*num_threads=*/1); + } + + __syncthreads(); + + pipeline_load_q.init_masks(ClusterShape{}); + pipeline_load_k.init_masks(ClusterShape{}); + pipeline_load_v.init_masks(ClusterShape{}); + pipeline_mma_s0.init_masks(ClusterShape{}); + pipeline_mma_s1.init_masks(ClusterShape{}); + pipeline_mma_corr.init_masks(ClusterShape{}); + + typename CollectiveMainloop::PipelineQ::PipelineState pipeline_load_q_consumer_state; + typename CollectiveMainloop::PipelineQ::PipelineState pipeline_load_q_producer_state = + cutlass::make_producer_start_state(); + + typename CollectiveMainloop::PipelineK::PipelineState pipeline_load_k_consumer_state; + typename CollectiveMainloop::PipelineK::PipelineState pipeline_load_k_producer_state = + cutlass::make_producer_start_state(); + + typename CollectiveMainloop::PipelineV::PipelineState pipeline_load_v_consumer_state; + typename CollectiveMainloop::PipelineV::PipelineState pipeline_load_v_producer_state = + cutlass::make_producer_start_state(); + + typename CollectiveMainloop::PipelineS::PipelineState pipeline_mma_s0_consumer_state; + typename CollectiveMainloop::PipelineS::PipelineState pipeline_mma_s0_producer_state = + cutlass::make_producer_start_state(); + + typename CollectiveMainloop::PipelineS::PipelineState pipeline_mma_s1_consumer_state; + typename CollectiveMainloop::PipelineS::PipelineState pipeline_mma_s1_producer_state = + cutlass::make_producer_start_state(); + + typename CollectiveMainloop::PipelineC::PipelineState pipeline_s0_corr_consumer_state; + typename CollectiveMainloop::PipelineC::PipelineState pipeline_s0_corr_producer_state = + cutlass::make_producer_start_state(); + + typename CollectiveMainloop::PipelineC::PipelineState pipeline_s1_corr_consumer_state; + typename CollectiveMainloop::PipelineC::PipelineState pipeline_s1_corr_producer_state = + cutlass::make_producer_start_state(); + + typename CollectiveMainloop::PipelineE::PipelineState pipeline_corr_epi_consumer_state; + typename CollectiveMainloop::PipelineE::PipelineState pipeline_corr_epi_producer_state = + cutlass::make_producer_start_state(); + + typename CollectiveMainloop::PipelineO::PipelineState pipeline_mma_corr_consumer_state; + typename CollectiveMainloop::PipelineO::PipelineState pipeline_mma_corr_producer_state = + cutlass::make_producer_start_state(); + + CollectiveMainloop mainloop; + CollectiveEpilogue epilogue{params.epilogue}; + + if (role == WarpRole::Softmax0 || role == WarpRole::Softmax1) { + warpgroup_reg_set(); + + CUTLASS_PRAGMA_NO_UNROLL + for (; tile_scheduler.is_valid(); ++tile_scheduler) { + auto blk_coord = tile_scheduler.get_block_coord(); + + auto logical_problem_shape = + apply_batch(params, params.problem_shape, get<2, 1>(blk_coord)); + + if (get<0>(blk_coord) * get<0>(TileShape{}) >= get<0>(logical_problem_shape)) { + continue; + } + + if (get<1>(logical_problem_shape) == 0) { // kv_len == 0 + continue; + } + + bool is_softmax_0 = role == WarpRole::Softmax0; + + mainloop.softmax( + is_softmax_0 ? 0 : 1, blk_coord, params.mainloop, logical_problem_shape, + is_softmax_0 ? pipeline_mma_s0 : pipeline_mma_s1, + is_softmax_0 ? pipeline_mma_s0_consumer_state : pipeline_mma_s1_consumer_state, + is_softmax_0 ? pipeline_s0_corr : pipeline_s1_corr, + is_softmax_0 ? pipeline_s0_corr_producer_state : pipeline_s1_corr_producer_state, + order_s01); + } + } else if (role == WarpRole::Correction) { + cutlass::arch::warpgroup_reg_dealloc(); + + CUTLASS_PRAGMA_NO_UNROLL + for (; tile_scheduler.is_valid(); ++tile_scheduler) { + auto blk_coord = tile_scheduler.get_block_coord(); + + auto logical_problem_shape = + apply_batch(params, params.problem_shape, get<2, 1>(blk_coord)); + + if (get<0>(blk_coord) * get<0>(TileShape{}) >= get<0>(logical_problem_shape)) { + continue; + } + + if (get<1>(logical_problem_shape) == 0) { // kv_len == 0 + mainloop.correction_empty(blk_coord, params.mainloop, logical_problem_shape, + params.problem_shape, shared_storage.epilogue, + pipeline_corr_epi, pipeline_corr_epi_producer_state, epilogue); + continue; + } + + mainloop.correction(blk_coord, params.mainloop, params.problem_shape, logical_problem_shape, + shared_storage.epilogue, pipeline_s0_corr, + pipeline_s0_corr_consumer_state, pipeline_s1_corr, + pipeline_s1_corr_consumer_state, pipeline_mma_corr, + pipeline_mma_corr_consumer_state, pipeline_corr_epi, + pipeline_corr_epi_producer_state, epilogue); + } + + if constexpr (NumWarpsEpilogue == 0) { + static_assert(NumWarpsCorrection == 1); + + uint32_t free_stage_ptr = shared_storage.tmem_base_ptr; + tmem_allocator.free(free_stage_ptr, TmemAllocator::Sm100TmemCapacityColumns); + } + + } else if (role == WarpRole::MMA) { + warpgroup_reg_set(); + + tmem_allocator.allocate(TmemAllocator::Sm100TmemCapacityColumns, + &shared_storage.tmem_base_ptr); + __syncwarp(); + + CUTLASS_PRAGMA_NO_UNROLL + for (; tile_scheduler.is_valid(); ++tile_scheduler) { + auto blk_coord = tile_scheduler.get_block_coord(); + + auto logical_problem_shape = + apply_batch(params, params.problem_shape, get<2, 1>(blk_coord)); + + if (get<0>(blk_coord) * get<0>(TileShape{}) >= get<0>(logical_problem_shape)) { + continue; + } + + if (get<1>(logical_problem_shape) == 0) { // kv_len == 0 + continue; + } + + mainloop.mma( + blk_coord, params.mainloop, logical_problem_shape, shared_storage.mainloop, + pipeline_load_q, pipeline_load_q_consumer_state, pipeline_load_k, + pipeline_load_k_consumer_state, pipeline_load_v, pipeline_load_v_consumer_state, + pipeline_mma_s0, pipeline_mma_s0_producer_state, pipeline_mma_s1, + pipeline_mma_s1_producer_state, pipeline_mma_corr, pipeline_mma_corr_producer_state); + } + } else if (role == WarpRole::Load) { + warpgroup_reg_set(); + + int work_idx = 0; + + CUTLASS_PRAGMA_NO_UNROLL + for (; tile_scheduler.is_valid(); ++tile_scheduler) { + shared_storage.barrier_O.wait((work_idx + 1) % 2); + auto blk_coord = tile_scheduler.get_block_coord(); + + auto logical_problem_shape = + apply_batch(params, params.problem_shape, get<2, 1>(blk_coord)); + + if (get<0>(blk_coord) * get<0>(TileShape{}) >= get<0>(logical_problem_shape)) { + continue; + } + + if (get<1>(logical_problem_shape) == 0) { // kv_len == 0 + work_idx++; + continue; + } + + mainloop.load(blk_coord, logical_problem_shape, params.mainloop, params.problem_shape, + shared_storage.mainloop, pipeline_load_q, pipeline_load_q_producer_state, + pipeline_load_k, pipeline_load_k_producer_state, pipeline_load_v, + pipeline_load_v_producer_state); + + work_idx++; + } + } else if (role == WarpRole::Epilogue) { + warpgroup_reg_set(); + + int work_idx = 0; + CUTLASS_PRAGMA_NO_UNROLL + for (; tile_scheduler.is_valid(); ++tile_scheduler) { + if (work_idx != 0) { + if (lane_predicate) { + shared_storage.barrier_O.arrive(0, lane_predicate); + } + } + auto blk_coord = tile_scheduler.get_block_coord(); + + auto logical_problem_shape = + apply_batch(params, params.problem_shape, get<2, 1>(blk_coord)); + + if (get<0>(blk_coord) * get<0>(TileShape{}) >= get<0>(logical_problem_shape)) { + continue; + } + + epilogue.store(blk_coord, logical_problem_shape, params.epilogue, params.problem_shape, + shared_storage.epilogue, pipeline_corr_epi, + pipeline_corr_epi_consumer_state); + + work_idx++; + } + + static_assert(NumWarpsEpilogue <= 1); + if constexpr (NumWarpsEpilogue == 1) { + uint32_t free_stage_ptr = shared_storage.tmem_base_ptr; + tmem_allocator.free(free_stage_ptr, TmemAllocator::Sm100TmemCapacityColumns); + } + + } else if (role == WarpRole::Empty) { + warpgroup_reg_set(); + + /* no-op, donate regs and exit */ + } + } +}; + +} // namespace cutlass::fmha::kernel diff --git a/include/flashinfer/attention/blackwell/kernel/sm100_fmha_gen_kernel_warpspecialized.hpp b/include/flashinfer/attention/blackwell/kernel/sm100_fmha_gen_kernel_warpspecialized.hpp new file mode 100644 index 0000000000000000000000000000000000000000..1e9a77502dfe89a342cab17b0a9c71132cacf860 --- /dev/null +++ b/include/flashinfer/attention/blackwell/kernel/sm100_fmha_gen_kernel_warpspecialized.hpp @@ -0,0 +1,530 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +#include "collective/fmha_fusion.hpp" +#include "cute/arch/tmem_allocator_sm100.hpp" +#include "cute/layout.hpp" +#include "cutlass/arch/arch.h" +#include "cutlass/cutlass.h" +#include "cutlass/kernel_hardware_info.h" +#include "cutlass/pipeline/pipeline.hpp" +#include "kernel/fmha_options.hpp" +#include "kernel/fmha_tile_scheduler.hpp" + +namespace cutlass::fmha::kernel { + +using namespace cute; +using namespace cutlass::fmha::collective; + +struct Sm100FmhaGenKernelWarpspecializedSchedule { + enum class WarpRole { Softmax0, Softmax1, Correction, MMA, Load, Epilogue, Empty }; + + static constexpr WarpRole warp_idx_to_WarpRole(int warp_idx) { + if (warp_idx == 0) return WarpRole::Softmax0; // 0 - 3 + if (warp_idx == 1) return WarpRole::MMA; // 12 + if (warp_idx == 2 || warp_idx == 3) return WarpRole::Load; // 13 + if (warp_idx == 4) return WarpRole::Softmax1; // 4 - 7 + if (warp_idx == 8) return WarpRole::Correction; // 8 - 11 + return WarpRole::Empty; // 15 + } + + static const int NumWarpsSoftmax = 1; + static const int NumWarpsCorrection = 1; + static const int NumWarpsEpilogue = 0; + static const int NumWarpsLoad = 2; + + static const int NumRegsSoftmax = 192; + static const int NumRegsCorrection = 104; + static const int NumRegsOther = 248; + static const int NumRegsEmpty = 24; + + static const int NumWarps = 12; +}; + +template +struct Sm100FmhaGenKernelWarpspecialized { + using TileShape = typename CollectiveMainloop::TileShape; + using ProblemShape = decltype(replace<0>(ProblemShapeIn{}, 0)); + + using WarpRole = typename KernelSchedule::WarpRole; + + constexpr WarpRole warp_idx_to_WarpRole(int warp_idx) { + return KernelSchedule::warp_idx_to_WarpRole(warp_idx); + } + + static const int NumWarpsSoftmax = KernelSchedule::NumWarpsSoftmax; + static const int NumWarpsCorrection = KernelSchedule::NumWarpsCorrection; + static const int NumWarpsEpilogue = KernelSchedule::NumWarpsEpilogue; + static const int NumWarpsLoad = KernelSchedule::NumWarpsLoad; + + static const int NumRegsSoftmax = KernelSchedule::NumRegsSoftmax; + static const int NumRegsCorrection = KernelSchedule::NumRegsCorrection; + static const int NumRegsOther = KernelSchedule::NumRegsOther; + static const int NumRegsEmpty = 24; + + static const int NumWarps = KernelSchedule::NumWarps; + + using ClusterShape = typename CollectiveMainloop::ClusterShape; + + using TmemAllocator = cute::TMEM::Allocator1Sm; + + struct SharedStorage { + typename CollectiveMainloop::TensorStorage mainloop; + typename CollectiveEpilogue::TensorStorage epilogue; + + struct PipelineStorage { + alignas(16) typename CollectiveMainloop::PipelineQ::SharedStorage load_q; + alignas(16) typename CollectiveMainloop::PipelineKV::SharedStorage load_kv; + alignas(16) typename CollectiveMainloop::PipelineS::SharedStorage mma_s0; + alignas(16) typename CollectiveMainloop::PipelineS::SharedStorage mma_s1; + alignas(16) typename CollectiveMainloop::PipelineC::SharedStorage s0_corr; + alignas(16) typename CollectiveMainloop::PipelineC::SharedStorage s1_corr; + alignas(16) typename CollectiveMainloop::PipelineO::SharedStorage mma_corr; + alignas(16) typename CollectiveMainloop::PipelineE::SharedStorage corr_epi; + alignas(16) typename CollectiveMainloop::OrderBarrierSoftmax::SharedStorage order_s01; + } pipelines; + + uint32_t tmem_base_ptr; + }; + + static constexpr int SharedStorageSize = sizeof(SharedStorage); + + using StrideQOrig = typename CollectiveMainloop::StrideQOrig; + using StrideOOrig = typename CollectiveMainloop::StrideOOrig; + using StrideQ = typename CollectiveMainloop::StrideQ; + using StrideO = typename CollectiveMainloop::StrideO; + using StrideCacheK = typename CollectiveMainloop::StrideCacheK; + using StrideCacheV = typename CollectiveMainloop::StrideCacheV; + using StrideNewK = typename CollectiveMainloop::StrideNewK; + using StrideNewV = typename CollectiveMainloop::StrideNewV; + using Element = typename CollectiveMainloop::Element; + using ElementAcc = typename CollectiveMainloop::ElementAcc; + using ElementOut = typename CollectiveMainloop::ElementOut; + + struct Arguments { + // _1, max_seqlen_k, head_dim, ((h_g, h_kv), b) + ProblemShapeIn problem_shape; + const int* seqlen_kv; + const int* cache_batch_idx; + + const Element* ptr_q; // 1 x D x (H x B) + StrideQOrig dQ; + const Element* ptr_new_k; // 1 x D x (H x B) + StrideNewK dNewK; + const Element* ptr_new_v; // 1 x D x (H x B) + StrideNewV dNewV; + + Element* ptr_cache_k; // seqlen_max x D x (H x B) + StrideCacheK dCacheK; + Element* ptr_cache_v; // seqlen_max x D x (H x B) + StrideCacheV dCacheV; + ElementOut* ptr_o; // 1 x D x (H x B) + StrideOOrig dO; + + cutlass::KernelHardwareInfo hw_info; + + ElementAcc scale_softmax = 0.0f; + }; + + struct Params { + ProblemShape problem_shape; + const int* seqlen_kv; + typename CollectiveMainloop::Params mainloop; + typename CollectiveEpilogue::Params epilogue; + typename TileScheduler::Params tile_scheduler; + }; + + static const int MinBlocksPerMultiprocessor = 1; + static const int MaxThreadsPerBlock = NumWarps * cutlass::NumThreadsPerWarp; + using ArchTag = cutlass::arch::Sm100; + + static size_t get_workspace_size(Arguments const& args) { return 0; } + static cutlass::Status initialize_workspace(Arguments const&, void*, cudaStream_t) { + return cutlass::Status::kSuccess; + } + + static bool can_implement(Arguments const& args) { return true; } + + static dim3 get_grid_shape(Params const& params) { + return TileScheduler::get_grid_shape(params.tile_scheduler); + } + + static dim3 get_block_shape() { + dim3 block(MaxThreadsPerBlock, 1, 1); + return block; + } + + static Params to_underlying_arguments(Arguments const& args, void* workspace) { + ProblemShape problem_shape = + replace<0>(args.problem_shape, static_cast(get<0>(args.problem_shape))); + CUTE_STATIC_ASSERT_V(get<0>(args.problem_shape) == _1{}); + StrideQ dQ = replace<0>(args.dQ, 0); + StrideO dO = replace<0>(args.dO, 0); + get<0>(problem_shape) = get<3, 0, 0>(args.problem_shape); + get<3, 0, 0>(problem_shape) = 1; + get<0>(dQ) = get<2, 0, 0>(dQ); + get<0>(dO) = get<2, 0, 0>(dO); + + typename CollectiveMainloop::Arguments mainloop_args{{ + args.cache_batch_idx, + args.ptr_q, + dQ, + args.ptr_new_k, + args.dNewK, + args.ptr_new_v, + args.dNewV, + args.ptr_cache_k, + args.dCacheK, + args.ptr_cache_v, + args.dCacheV, + }, + args.scale_softmax}; + + typename CollectiveEpilogue::Arguments epilogue_args{ + args.ptr_o, + dO, + }; + + return Params{ + problem_shape, args.seqlen_kv, + CollectiveMainloop::to_underlying_arguments(problem_shape, mainloop_args, workspace), + CollectiveEpilogue::to_underlying_arguments(problem_shape, epilogue_args, workspace), + TileScheduler::to_underlying_arguments(problem_shape, args.hw_info, ClusterShape{}, + TileShape{})}; + } + + CUTLASS_DEVICE auto apply_batch(const Params& params, ProblemShape const& problem_shape, + int batch_idx) { + ProblemShape result = problem_shape; + get<1>(result) = params.seqlen_kv[batch_idx]; + if (params.mainloop.load.ptr_new_k != nullptr) { + get<1>(result) += 1; + } + return result; + } + + CUTLASS_DEVICE void operator()(const Params& params, char* smem) { + TileScheduler tile_scheduler{params.tile_scheduler}; + + int warp_idx = cutlass::canonical_warp_idx_sync(); + auto role = warp_idx_to_WarpRole(warp_idx); + uint32_t lane_predicate = cute::elect_one_sync(); + + if (role == WarpRole::Load && lane_predicate) { + CollectiveMainloop::prefetch_tma_descriptors(params.mainloop); + } + + if (role == WarpRole::Epilogue && lane_predicate) { + CollectiveEpilogue::prefetch_tma_descriptors(params.epilogue); + } + + SharedStorage& shared_storage = *reinterpret_cast(smem); + + typename CollectiveMainloop::PipelineQ::Params pipeline_load_q_params; + if (role == WarpRole::Load) { + pipeline_load_q_params.role = CollectiveMainloop::PipelineQ::ThreadCategory::Producer; + } + if (role == WarpRole::MMA) { + pipeline_load_q_params.role = CollectiveMainloop::PipelineQ::ThreadCategory::Consumer; + } + pipeline_load_q_params.producer_arv_count = NumWarpsLoad * cutlass::NumThreadsPerWarp; + typename CollectiveMainloop::PipelineQ pipeline_load_q( + shared_storage.pipelines.load_q, pipeline_load_q_params, ClusterShape{}, cute::true_type{}, + /*mask calc*/ cute::false_type{}); + + typename CollectiveMainloop::PipelineKV::Params pipeline_load_kv_params; + if (role == WarpRole::Load) { + pipeline_load_kv_params.role = CollectiveMainloop::PipelineKV::ThreadCategory::Producer; + } + if (role == WarpRole::MMA) { + pipeline_load_kv_params.role = CollectiveMainloop::PipelineKV::ThreadCategory::Consumer; + } + pipeline_load_kv_params.producer_arv_count = NumWarpsLoad * cutlass::NumThreadsPerWarp; + typename CollectiveMainloop::PipelineKV pipeline_load_kv( + shared_storage.pipelines.load_kv, pipeline_load_kv_params, ClusterShape{}, + /*barrier init*/ cute::true_type{}, /*mask calc*/ cute::false_type{}); + + typename CollectiveMainloop::PipelineS::Params pipeline_mma_s0_params; + if (role == WarpRole::MMA) { + pipeline_mma_s0_params.role = CollectiveMainloop::PipelineS::ThreadCategory::Producer; + } + if (role == WarpRole::Softmax0) { + pipeline_mma_s0_params.role = CollectiveMainloop::PipelineS::ThreadCategory::Consumer; + } + pipeline_mma_s0_params.consumer_arv_count = NumWarpsSoftmax * cutlass::NumThreadsPerWarp; + typename CollectiveMainloop::PipelineS pipeline_mma_s0( + shared_storage.pipelines.mma_s0, pipeline_mma_s0_params, ClusterShape{}, + /*barrier init*/ cute::true_type{}, /*mask calc*/ cute::false_type{}); + + typename CollectiveMainloop::PipelineS::Params pipeline_mma_s1_params; + if (role == WarpRole::MMA) { + pipeline_mma_s1_params.role = CollectiveMainloop::PipelineS::ThreadCategory::Producer; + } + if (role == WarpRole::Softmax1) { + pipeline_mma_s1_params.role = CollectiveMainloop::PipelineS::ThreadCategory::Consumer; + } + pipeline_mma_s1_params.consumer_arv_count = NumWarpsSoftmax * cutlass::NumThreadsPerWarp; + typename CollectiveMainloop::PipelineS pipeline_mma_s1( + shared_storage.pipelines.mma_s1, pipeline_mma_s1_params, ClusterShape{}, + /*barrier init*/ cute::true_type{}, /*mask calc*/ cute::false_type{}); + + typename CollectiveMainloop::PipelineC::Params pipeline_s0_corr_params; + if (role == WarpRole::Softmax0) { + pipeline_s0_corr_params.role = CollectiveMainloop::PipelineC::ThreadCategory::Producer; + } + if (role == WarpRole::Correction) { + pipeline_s0_corr_params.role = CollectiveMainloop::PipelineC::ThreadCategory::Consumer; + } + pipeline_s0_corr_params.producer_arv_count = NumWarpsSoftmax * cutlass::NumThreadsPerWarp; + pipeline_s0_corr_params.consumer_arv_count = NumWarpsCorrection * cutlass::NumThreadsPerWarp; + typename CollectiveMainloop::PipelineC pipeline_s0_corr(shared_storage.pipelines.s0_corr, + pipeline_s0_corr_params, + /*barrier init*/ cute::true_type{}); + + typename CollectiveMainloop::PipelineC::Params pipeline_s1_corr_params; + if (role == WarpRole::Softmax1) { + pipeline_s1_corr_params.role = CollectiveMainloop::PipelineC::ThreadCategory::Producer; + } + if (role == WarpRole::Correction) { + pipeline_s1_corr_params.role = CollectiveMainloop::PipelineC::ThreadCategory::Consumer; + } + pipeline_s1_corr_params.producer_arv_count = NumWarpsSoftmax * cutlass::NumThreadsPerWarp; + pipeline_s1_corr_params.consumer_arv_count = NumWarpsCorrection * cutlass::NumThreadsPerWarp; + typename CollectiveMainloop::PipelineC pipeline_s1_corr(shared_storage.pipelines.s1_corr, + pipeline_s1_corr_params, + /*barrier init*/ cute::true_type{}); + + typename CollectiveMainloop::PipelineO::Params pipeline_mma_corr_params; + if (role == WarpRole::MMA) { + pipeline_mma_corr_params.role = CollectiveMainloop::PipelineO::ThreadCategory::Producer; + } + if (role == WarpRole::Correction) { + pipeline_mma_corr_params.role = CollectiveMainloop::PipelineO::ThreadCategory::Consumer; + } + pipeline_mma_corr_params.consumer_arv_count = NumWarpsCorrection * cutlass::NumThreadsPerWarp; + typename CollectiveMainloop::PipelineO pipeline_mma_corr( + shared_storage.pipelines.mma_corr, pipeline_mma_corr_params, ClusterShape{}, + /*barrier init*/ cute::true_type{}, /*mask calc*/ cute::false_type{}); + + typename CollectiveMainloop::PipelineE::Params pipeline_corr_epi_params; + if (role == WarpRole::Correction) { + pipeline_corr_epi_params.role = CollectiveMainloop::PipelineE::ThreadCategory::Producer; + } + if (role == WarpRole::Epilogue) { + pipeline_corr_epi_params.role = CollectiveMainloop::PipelineE::ThreadCategory::Consumer; + } + pipeline_corr_epi_params.producer_arv_count = NumWarpsCorrection * cutlass::NumThreadsPerWarp; + pipeline_corr_epi_params.consumer_arv_count = NumWarpsEpilogue * cutlass::NumThreadsPerWarp; + typename CollectiveMainloop::PipelineE pipeline_corr_epi(shared_storage.pipelines.corr_epi, + pipeline_corr_epi_params, + /*barrier init*/ cute::true_type{}); + + typename CollectiveMainloop::OrderBarrierSoftmax::Params params_order_s01; + params_order_s01.group_id = role == WarpRole::Softmax1 ? 1 : 0; + params_order_s01.group_size = NumWarpsSoftmax * cutlass::NumThreadsPerWarp; + typename CollectiveMainloop::OrderBarrierSoftmax order_s01(shared_storage.pipelines.order_s01, + params_order_s01); + + TmemAllocator tmem_allocator; + + __syncthreads(); + + pipeline_load_q.init_masks(ClusterShape{}); + pipeline_load_kv.init_masks(ClusterShape{}); + pipeline_mma_s0.init_masks(ClusterShape{}); + pipeline_mma_s1.init_masks(ClusterShape{}); + pipeline_mma_corr.init_masks(ClusterShape{}); + + typename CollectiveMainloop::PipelineQ::PipelineState pipeline_load_q_consumer_state; + typename CollectiveMainloop::PipelineQ::PipelineState pipeline_load_q_producer_state = + cutlass::make_producer_start_state(); + + typename CollectiveMainloop::PipelineKV::PipelineState pipeline_load_kv_consumer_state; + typename CollectiveMainloop::PipelineKV::PipelineState pipeline_load_kv_producer_state = + cutlass::make_producer_start_state(); + + typename CollectiveMainloop::PipelineS::PipelineState pipeline_mma_s0_consumer_state; + typename CollectiveMainloop::PipelineS::PipelineState pipeline_mma_s0_producer_state = + cutlass::make_producer_start_state(); + + typename CollectiveMainloop::PipelineS::PipelineState pipeline_mma_s1_consumer_state; + typename CollectiveMainloop::PipelineS::PipelineState pipeline_mma_s1_producer_state = + cutlass::make_producer_start_state(); + + typename CollectiveMainloop::PipelineC::PipelineState pipeline_s0_corr_consumer_state; + typename CollectiveMainloop::PipelineC::PipelineState pipeline_s0_corr_producer_state = + cutlass::make_producer_start_state(); + + typename CollectiveMainloop::PipelineC::PipelineState pipeline_s1_corr_consumer_state; + typename CollectiveMainloop::PipelineC::PipelineState pipeline_s1_corr_producer_state = + cutlass::make_producer_start_state(); + + typename CollectiveMainloop::PipelineE::PipelineState pipeline_corr_epi_consumer_state; + typename CollectiveMainloop::PipelineE::PipelineState pipeline_corr_epi_producer_state = + cutlass::make_producer_start_state(); + + typename CollectiveMainloop::PipelineO::PipelineState pipeline_mma_corr_consumer_state; + typename CollectiveMainloop::PipelineO::PipelineState pipeline_mma_corr_producer_state = + cutlass::make_producer_start_state(); + + CollectiveMainloop mainloop; + CollectiveEpilogue epilogue(params.epilogue); + + if (role == WarpRole::Softmax0 || role == WarpRole::Softmax1) { + warpgroup_reg_set(); + + CUTLASS_PRAGMA_NO_UNROLL + for (; tile_scheduler.is_valid(); ++tile_scheduler) { + auto blk_coord = tile_scheduler.get_block_coord(); + + auto logical_problem_shape = + apply_batch(params, params.problem_shape, get<2, 1>(blk_coord)); + + if (get<0>(blk_coord) * get<0>(TileShape{}) >= get<0>(logical_problem_shape)) { + continue; + } + + bool is_softmax_0 = role == WarpRole::Softmax0; + + mainloop.softmax( + is_softmax_0 ? 0 : 1, blk_coord, params.mainloop, logical_problem_shape, + is_softmax_0 ? pipeline_mma_s0 : pipeline_mma_s1, + is_softmax_0 ? pipeline_mma_s0_consumer_state : pipeline_mma_s1_consumer_state, + is_softmax_0 ? pipeline_s0_corr : pipeline_s1_corr, + is_softmax_0 ? pipeline_s0_corr_producer_state : pipeline_s1_corr_producer_state, + order_s01); + } + } else if (role == WarpRole::Correction) { + cutlass::arch::warpgroup_reg_dealloc(); + + CUTLASS_PRAGMA_NO_UNROLL + for (; tile_scheduler.is_valid(); ++tile_scheduler) { + auto blk_coord = tile_scheduler.get_block_coord(); + + auto logical_problem_shape = + apply_batch(params, params.problem_shape, get<2, 1>(blk_coord)); + + if (get<0>(blk_coord) * get<0>(TileShape{}) >= get<0>(logical_problem_shape)) { + continue; + } + + mainloop.correction( + blk_coord, params.mainloop, logical_problem_shape, shared_storage.epilogue, + pipeline_s0_corr, pipeline_s0_corr_consumer_state, pipeline_s1_corr, + pipeline_s1_corr_consumer_state, pipeline_mma_corr, pipeline_mma_corr_consumer_state, + pipeline_corr_epi, pipeline_corr_epi_producer_state, epilogue); + } + + if constexpr (NumWarpsEpilogue == 0) { + static_assert(NumWarpsCorrection == 1); + + uint32_t free_stage_ptr = shared_storage.tmem_base_ptr; + tmem_allocator.free(free_stage_ptr, TmemAllocator::Sm100TmemCapacityColumns); + } + + } else if (role == WarpRole::MMA) { + warpgroup_reg_set(); + + tmem_allocator.allocate(TmemAllocator::Sm100TmemCapacityColumns, + &shared_storage.tmem_base_ptr); + __syncwarp(); + + CUTLASS_PRAGMA_NO_UNROLL + for (; tile_scheduler.is_valid(); ++tile_scheduler) { + auto blk_coord = tile_scheduler.get_block_coord(); + + auto logical_problem_shape = + apply_batch(params, params.problem_shape, get<2, 1>(blk_coord)); + + if (get<0>(blk_coord) * get<0>(TileShape{}) >= get<0>(logical_problem_shape)) { + continue; + } + + mainloop.mma(blk_coord, params.mainloop, logical_problem_shape, shared_storage.mainloop, + pipeline_load_q, pipeline_load_q_consumer_state, pipeline_load_kv, + pipeline_load_kv_consumer_state, pipeline_mma_s0, + pipeline_mma_s0_producer_state, pipeline_mma_s1, + pipeline_mma_s1_producer_state, pipeline_mma_corr, + pipeline_mma_corr_producer_state); + } + } else if (role == WarpRole::Load) { + warpgroup_reg_set(); + + CUTLASS_PRAGMA_NO_UNROLL + for (; tile_scheduler.is_valid(); ++tile_scheduler) { + auto blk_coord = tile_scheduler.get_block_coord(); + + auto logical_problem_shape = + apply_batch(params, params.problem_shape, get<2, 1>(blk_coord)); + + if (get<0>(blk_coord) * get<0>(TileShape{}) >= get<0>(logical_problem_shape)) { + continue; + } + + mainloop.load(blk_coord, logical_problem_shape, params.mainloop, params.problem_shape, + shared_storage.mainloop, pipeline_load_q, pipeline_load_q_producer_state, + pipeline_load_kv, pipeline_load_kv_producer_state); + } + } else if (role == WarpRole::Epilogue) { + warpgroup_reg_set(); + + CUTLASS_PRAGMA_NO_UNROLL + for (; tile_scheduler.is_valid(); ++tile_scheduler) { + auto blk_coord = tile_scheduler.get_block_coord(); + + auto logical_problem_shape = + apply_batch(params, params.problem_shape, get<2, 1>(blk_coord)); + + if (get<0>(blk_coord) * get<0>(TileShape{}) >= get<0>(logical_problem_shape)) { + continue; + } + + epilogue.store(blk_coord, logical_problem_shape, params.epilogue, params.problem_shape, + shared_storage.epilogue, pipeline_corr_epi, + pipeline_corr_epi_consumer_state); + } + + static_assert(NumWarpsEpilogue <= 1); + if constexpr (NumWarpsEpilogue == 1) { + uint32_t free_stage_ptr = shared_storage.tmem_base_ptr; + tmem_allocator.free(free_stage_ptr, TmemAllocator::Sm100TmemCapacityColumns); + } + + } else if (role == WarpRole::Empty) { + warpgroup_reg_set(); + + /* no-op, donate regs and exit */ + } + } +}; + +} // namespace cutlass::fmha::kernel diff --git a/include/flashinfer/attention/blackwell/kernel/sm100_fmha_mla_reduction.hpp b/include/flashinfer/attention/blackwell/kernel/sm100_fmha_mla_reduction.hpp new file mode 100644 index 0000000000000000000000000000000000000000..7f12313ded953afdb356b1feedff986a28f35be6 --- /dev/null +++ b/include/flashinfer/attention/blackwell/kernel/sm100_fmha_mla_reduction.hpp @@ -0,0 +1,195 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +#pragma once + +#include "cute/tensor.hpp" +#include "cutlass/arch/arch.h" +#include "cutlass/cutlass.h" + +namespace cutlass::fmha::kernel { + +using namespace cute; +template +struct Sm100FmhaMlaReductionKernel { + static const int SharedStorageSize = 0; + static const int MaxThreadsPerBlock = 128; + static const int MinBlocksPerMultiprocessor = 1; + + using ArchTag = cutlass::arch::Sm100; + + static_assert(kHeadDimLatent % MaxThreadsPerBlock == 0); + struct Arguments { + ElementAcc* ptr_oaccum = nullptr; + ElementOut* ptr_o = nullptr; + ElementAcc* ptr_lseaccum = nullptr; + ElementAcc* ptr_lse = nullptr; + ElementScale scale = 1.f; + int num_batches = 0; + int split_kv = -1; + int dim_k = -1; + int* ptr_seq = nullptr; + int* ptr_split_kv = nullptr; + int tile_shape_s = 128; + }; + using Params = Arguments; + + static Params to_underlying_arguments(Arguments const& args, void* workspace) { + return {args.ptr_oaccum, args.ptr_o, args.ptr_lseaccum, args.ptr_lse, + args.scale, args.num_batches, args.split_kv, args.dim_k, + args.ptr_seq, args.ptr_split_kv, args.tile_shape_s}; + } + + static size_t get_workspace_size(Arguments const& /*args*/) { return 0; } + + static Status initialize_workspace(Arguments const& /*args*/, void* /*ws*/, + cudaStream_t /*stream*/) { + return Status::kSuccess; + } + + static dim3 get_grid_shape(Params const& params) { + return dim3(kNumHeads, 1, params.num_batches); + } + + static dim3 get_block_shape() { return dim3(MaxThreadsPerBlock, 1, 1); } + + static bool can_implement(Arguments const& args) { + if (args.num_batches <= 0) return false; + if (args.split_kv <= 0) return false; + return true; + } + + CUTLASS_DEVICE void operator()(Params const& params, char* smem_raw) { + if (params.split_kv <= 1) return; + auto blk_coord = make_coord(blockIdx.x, _0{}, blockIdx.z); + + __shared__ ElementAcc sLseScale[kMaxSplits]; + const size_t offset_lseaccum = + get<0>(blk_coord) + kNumHeads * params.split_kv * get<2>(blk_coord); + const size_t offset_lse = get<0>(blk_coord) + kNumHeads * get<2>(blk_coord); + + Tensor gLSEaccum = make_tensor(make_gmem_ptr(params.ptr_lseaccum + offset_lseaccum), + make_shape(params.split_kv), Stride>{}); + + Tensor gLSE = + make_tensor(make_gmem_ptr(params.ptr_lse + offset_lse), Shape<_1>{}, Stride<_1>{}); + + auto dim_k = params.ptr_seq == nullptr ? params.dim_k : params.ptr_seq[get<2>(blk_coord)]; + auto local_split_kv = + params.ptr_split_kv == nullptr ? params.split_kv : params.ptr_split_kv[get<2>(blk_coord)]; + auto k_tile_total = ceil_div(dim_k, params.tile_shape_s); + auto k_tile_per_cta = ceil_div(k_tile_total, local_split_kv); + local_split_kv = ceil_div(k_tile_total, k_tile_per_cta); + + int warp_idx = cutlass::canonical_warp_idx_sync(); + if (warp_idx == 0) { + constexpr int kNLsePerThread = cute::ceil_div(kMaxSplits, 32); + + ElementAcc local_lse[kNLsePerThread]; + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < kNLsePerThread; ++i) { + const int split = i * 32 + threadIdx.x; + local_lse[i] = split < local_split_kv ? gLSEaccum(split) + : -std::numeric_limits::infinity(); + } + + ElementAcc lse_max = -std::numeric_limits::infinity(); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < kNLsePerThread; ++i) { + lse_max = max(lse_max, local_lse[i]); + } + CUTLASS_PRAGMA_UNROLL + for (int offset = 16; offset >= 1; offset /= 2) { + lse_max = max(lse_max, __shfl_xor_sync(0xffffffff, lse_max, offset)); + } + lse_max = lse_max == -std::numeric_limits::infinity() + ? 0.0f + : lse_max; // In case all local LSEs are -inf + lse_max = __shfl_sync(0xffffffff, lse_max, 0); + + ElementAcc sum_lse = 0; + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < kNLsePerThread; ++i) { + sum_lse = sum_lse + expf(local_lse[i] - lse_max); + } + + CUTLASS_PRAGMA_UNROLL + for (int offset = 16; offset >= 1; offset /= 2) { + sum_lse = sum_lse + __shfl_xor_sync(0xffffffff, sum_lse, offset); + } + + sum_lse = __shfl_sync(0xffffffff, sum_lse, 0); + + ElementAcc global_lse = (sum_lse == 0.f || sum_lse != sum_lse) + ? std::numeric_limits::infinity() + : logf(sum_lse) + lse_max; + if (threadIdx.x == 0 and params.ptr_lse != nullptr) { + gLSE(0) = global_lse; + } + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < kNLsePerThread; ++i) { + const int split = i * 32 + threadIdx.x; + if (split < local_split_kv) { + sLseScale[split] = expf(local_lse[i] - global_lse); + } + } + } + __syncthreads(); + + constexpr int Elements = kHeadDimLatent / MaxThreadsPerBlock; + const size_t offset_oaccum = + kHeadDimLatent * params.split_kv * (get<0>(blk_coord) + kNumHeads * get<2>(blk_coord)); + Tensor gOaccum = make_tensor(make_gmem_ptr(params.ptr_oaccum + offset_oaccum), + Shape>{}, Stride<_1>{}); + ElementAcc local_val[Elements] = {0}; + for (int split = 0; split < local_split_kv; ++split) { + ElementAcc lse_scale = sLseScale[split]; + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < Elements; ++i) { + local_val[i] += lse_scale * gOaccum(threadIdx.x + MaxThreadsPerBlock * i); + } + gOaccum.data() = gOaccum.data() + kHeadDimLatent; + } + auto ptr_o_local = + params.ptr_o + (get<0>(blk_coord) + get<2>(blk_coord) * kNumHeads) * kHeadDimLatent; + Tensor gO = make_tensor(make_gmem_ptr(ptr_o_local), Shape>{}, Stride<_1>{}); + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < Elements; ++i) { + gO(threadIdx.x + MaxThreadsPerBlock * i) = static_cast(local_val[i]); + } + } +}; + +} // namespace cutlass::fmha::kernel diff --git a/include/flashinfer/attention/blackwell/kernel/sm100_fmha_mla_tma_warpspecialized.hpp b/include/flashinfer/attention/blackwell/kernel/sm100_fmha_mla_tma_warpspecialized.hpp new file mode 100644 index 0000000000000000000000000000000000000000..66f31c3e25118ee93353ca338a86c931ee0352e6 --- /dev/null +++ b/include/flashinfer/attention/blackwell/kernel/sm100_fmha_mla_tma_warpspecialized.hpp @@ -0,0 +1,1941 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +#pragma once + +#include "../common/pow_2.hpp" +#include "cute/arch/simd_sm100.hpp" +#include "cute/tensor.hpp" +#include "cutlass/arch/arch.h" +#include "cutlass/arch/memory_sm80.h" +#include "cutlass/cutlass.h" +#include "cutlass/epilogue/thread/linear_combination.h" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "gather_tensor.hpp" // from examples/common + +namespace cutlass::fmha::kernel { + +using namespace cute; + +template +struct Sm100FmhaMlaKernelTmaWarpspecialized { + + using Element = Element_; + using ElementAcc = ElementAcc_; + using ElementOut = ElementOut_; + using ElementLSE = ElementLSE_; + + // only 2Sm mode is supported + static const bool kIs2Sm = true; + static const int MaxThreadsPerBlock = 256; + static const int MinBlocksPerMultiprocessor = 1; + static const int TotalSNum = 2; + static const int TotalPNum = 2; + using ArchTag = cutlass::arch::Sm100; + + using ClusterShape = cute::conditional_t, Shape<_1, _1, _1>>; + + using TileShapeH = tuple_element_t<0, TileShape>; + using TileShapeS = tuple_element_t<1, TileShape>; + using TileShapeD = tuple_element_t<2, TileShape>; + + using TileShapeL = tuple_element_t<0, TileShapeD>; + using TileShapeR = tuple_element_t<1, TileShapeD>; + static_assert(TileShapeL{} % TileShapeR{} == 0, "Rope head dim must divide latent head dim"); + + using ProblemShape = Shape; + using TensorStride = Stride; + using TmemAllocator = + cute::conditional_t; + + static_assert(TileShapeH{} == 128); + static const int kWarpsInN = kIs2Sm ? 2 : 1; + + static const int kNumComputeWarps = 4; + static const int kNumLoadWarps = kIsCpAsync ? 2 : 1; + + enum class WarpRole { + kMma = 0x1, + kLoad = 0x2, + kCompute = 0x3, + kLoadPageTable = 0x4, + kEmpty = 0x0 + }; + + static const long long unsigned int kWarpAssignment = + kIsCpAsync ? 0x4221'3333ull : 0x0021'3333ull; + + static CUTLASS_DEVICE WarpRole warp_idx_to_role(int warp_idx) { + return static_cast((kWarpAssignment >> (4 * warp_idx)) & 0xF); + } + + static const int Alignment = 128 / sizeof_bits_v; + static const int AlignmentOut = 128 / sizeof_bits_v; + + using TileShapeQK = Shape; + static const int StagesQK = 24 / sizeof(Element); // free parameter + static const int IterationsQKLatent = decltype(TileShapeL{} / get<2>(TileShapeQK{}))::value; + static const int IterationsQKRope = decltype(TileShapeR{} / get<2>(TileShapeQK{}))::value; + static const int IterationsQK = IterationsQKLatent + IterationsQKRope; + + using Schedule = cute::conditional_t; + using CollectiveMmaQK = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, Element, TensorStride, Alignment, + Element, TensorStride, Alignment, ElementAcc, TileShapeQK, ClusterShape, + cutlass::gemm::collective::StageCount, Schedule>::CollectiveOp; + using TiledMmaQK = typename CollectiveMmaQK::TiledMma; + using CtaShapeQK = typename CollectiveMmaQK::CtaShape_MNK; + + // chosen for unified smem staging between K and V + using TileShapePV = Shape; + using TransposeTensorStride = decltype(select<1, 0, 2>(TensorStride{})); + static const int StagesPV = StagesQK; // not sure why, but must be at least two. check pipes + static const int IterationsPV_K = decltype(TileShapeS{} / get<2>(TileShapePV{}))::value; + static const int IterationsPV_N = decltype(TileShapeL{} / get<1>(TileShapePV{}))::value; + + using CollectiveMmaPV = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, Element, TensorStride, Alignment, + Element, TransposeTensorStride, Alignment, ElementAcc, TileShapePV, ClusterShape, + cutlass::gemm::collective::StageCount, Schedule>::CollectiveOp; + using CtaShapePV = typename CollectiveMmaPV::CtaShape_MNK; + static_assert(std::is_same_v); + + using TiledMmaPV = typename CollectiveMmaPV::TiledMma; + + using AtomThrShapeMNK = typename CollectiveMmaQK::AtomThrShapeMNK; + static_assert(typename CollectiveMmaQK::AtomThrShapeMNK{} == + typename CollectiveMmaPV::AtomThrShapeMNK{}, + "schedule must match"); + + static const int StagesPageTable = kIsCpAsync ? StagesPV : 1; + + // pipelines from load to mma, PipelineTmaUmmaAsync, stages tbd + // use expect_tx for Q load + using PipelineLoadQK = + cute::conditional_t, + PipelineTmaUmmaAsync>; + using PipelineLoadPV = PipelineLoadQK; + // pipeline from mma (Q@K) to softmax, PipelineUmmaAsync, 2 stages + using PipelineS = PipelineUmmaAsync; + // pipeline from softmax (P) to mma (bmm2), PipelineUmmaAsync, 2 stages + using PipelineP = PipelineUmmaConsumerAsync; + // pipeline from mma to softmax (for rescale), PipelineUmmaAsync, 1 stage + using PipelineO = PipelineUmmaAsync<1, AtomThrShapeMNK>; + + using PipelinePT = PipelineAsync; + + struct PipelineStorage { + alignas(16) typename PipelineLoadQK::SharedStorage load_qk; + alignas(16) typename PipelineS::SharedStorage mma_s; + alignas(16) typename PipelineP::SharedStorage p_mma; + alignas(16) typename PipelineO::SharedStorage mma_o; + alignas(16) typename PipelinePT::SharedStorage load_page_table; + }; + + template + static CUTE_DEVICE constexpr auto unstageSmemLayout(Layout const& layout, Stages stages = {}) { + return composition(layout, make_tuple(_, _, _, make_layout(stages))); + } + + using SmemLayoutQ = + decltype(unstageSmemLayout(typename CollectiveMmaQK::SmemLayoutA{}, Int{})); + using SmemLayoutKC = typename CollectiveMmaQK::SmemLayoutB; + using SmemLayoutVC = typename CollectiveMmaPV::SmemLayoutB; + using SmemLayoutP = decltype(unstageSmemLayout(typename CollectiveMmaPV::SmemLayoutA{}, + make_shape(Int{}, _2{}))); + + static const int kBytesLoadQ = + size(AtomThrShapeMNK{}) * + cutlass::bits_to_bytes(cosize(take<0, 3>(SmemLayoutQ{})) * cute::sizeof_bits_v); + static const int kBytesLoadKC = + size(AtomThrShapeMNK{}) * + cutlass::bits_to_bytes(cosize(take<0, 3>(SmemLayoutKC{})) * cute::sizeof_bits_v); + static const int kBytesLoadVC = + size(AtomThrShapeMNK{}) * + cutlass::bits_to_bytes(cosize(take<0, 3>(SmemLayoutVC{})) * cute::sizeof_bits_v); + // pre-condition for overlapped smem staging + static_assert(kBytesLoadKC == kBytesLoadVC); + static_assert(StagesQK == StagesPV); + + static const int kTransactionsBytesLoadQK = kBytesLoadKC; + static const int kTransactionsBytesLoadExtraQ = kBytesLoadQ; + static const int kTransactionsBytesLoadPV = kBytesLoadVC; + + static const int kNamedBarrierExchange = + (int)cutlass::arch::ReservedNamedBarriers::TransformBarrier; + // This Named Barrier is introduced to solve Q tile loading overwritten issue when enable + // persistent tile scheduler for FP8 MLA. + static const int kNamedBarrierEpilogue = + (int)cutlass::arch::ReservedNamedBarriers::EpilogueBarrier; + // + static const int kNamedBarrierTmemDealloc = + (int)cutlass::arch::ReservedNamedBarriers::TmemAllocBarrier; + + enum class TmemAllocation : uint32_t { + kSizeS = TileShapeS::value / kWarpsInN, + // Overall + kSizeO = TileShapeL::value / kWarpsInN, + // Between accumulators we loop over + kSizeAccO = decltype(get<1>(TileShapePV{}))::value / kWarpsInN, + kNumS = TotalSNum, + kNumP = TotalPNum, + kNumO = 1, + kS0 = 0, + kS1 = kS0 + kSizeS, + kO0 = kS1 + kSizeS, + kTotal = kO0 + kSizeO + }; + + static_assert(static_cast(TmemAllocation::kTotal) <= TmemAllocator::Sm100TmemCapacityColumns, + "using too much tmem"); + + struct TensorStorage { + // to communicate max and row_sum + cute::array smem_exchange; + cute::array smem_page_table; + alignas(2048) cute::array> smem_q; + union { + alignas(2048) cute::array> smem_kc; + alignas(2048) cute::array> smem_vc; + }; + alignas(2048) cute::array> smem_p; + }; + + struct SharedStorage { + PipelineStorage pipelines; + TensorStorage tensors; + uint32_t tmem_base_ptr; + }; + + static const int SharedStorageSize = sizeof(SharedStorage); + static_assert(SharedStorageSize <= cutlass::arch::sm100_smem_capacity_bytes, + "using too much smem"); + + struct MainloopArguments { + ElementAcc softmax_scale; + + // all tensors strides are (num_heads or seqlen, head_dim, batch) + // head_dim stride is always 1 + Element* ptr_q_latent; + TensorStride stride_q_latent; + Element* ptr_q_rope; + TensorStride stride_q_rope; + + Element* ptr_c_latent; + TensorStride stride_c_latent; + Element* ptr_k_rope; + TensorStride stride_k_rope; + + // for paged attention, we interpret what was previously [batch, seqlen] + // as [page_count, page_size], and index according to page_table + int* ptr_seq = nullptr; + int* ptr_page_table = nullptr; + // page table is [batch, seqlen or similar] + Stride<_1, int> stride_page_table = {}; + int page_count = 0; + int page_size = TileShapeS{}; // powers of two if kIsCpAsync, otherwise TileShapeS + }; + + struct EpilogueArguments { + ElementOut* ptr_o = nullptr; + TensorStride stride_o; + ElementLSE* ptr_lse = nullptr; + Stride<_1, int> stride_lse; + ElementAcc output_scale = 1.0f; + }; + + struct Arguments { + // (num_heads=128, seqlen, (d_latent=512, d_rope=64), batch_count) + // for paged attention, seqlen is max seqlen + ProblemShape problem_shape; + MainloopArguments mainloop; + EpilogueArguments epilogue; + KernelHardwareInfo hw_info; + int split_kv = -1; + int* ptr_split_kv = nullptr; + }; + + using TmaLoadQLatent = typename CollectiveMmaQK::Params::TMA_A; + using TmaLoadQRope = typename CollectiveMmaQK::Params::TMA_A; + using TmaLoadCLatent = typename CollectiveMmaQK::Params::TMA_B; + using TmaLoadKRope = typename CollectiveMmaQK::Params::TMA_B; + using TmaLoadCLatentTranspose = typename CollectiveMmaPV::Params::TMA_B; + + struct MainloopParams { + TmaLoadQLatent tma_load_q_latent; + TmaLoadQRope tma_load_q_rope; + TmaLoadCLatent tma_load_c_latent; + TmaLoadKRope tma_load_k_rope; + TmaLoadCLatentTranspose tma_load_c_latent_transpose; + }; + + struct EpilogueParams { + ElementOut* ptr_o = nullptr; + ElementAcc* ptr_o_acc = nullptr; + TensorStride stride_o; + TensorStride stride_o_acc; + ElementLSE* ptr_lse = nullptr; + ElementLSE* ptr_lse_acc = nullptr; + Stride<_1, int> stride_lse; + Stride<_1, int> stride_lse_acc; + ElementAcc output_scale = 1.0f; + }; + + struct Params { + ProblemShape problem_shape; + MainloopArguments mainloop; + EpilogueParams epilogue; + MainloopParams mainloop_params; + typename TileScheduler::Params tile_scheduler; + int split_kv = -1; + int* ptr_split_kv = nullptr; + }; + + static Params to_underlying_arguments(Arguments const& args, void* workspace) { + // workspace = nullptr; // let's get an error if one of these needs workspace + + auto [H, K, D, B] = args.problem_shape; + auto [L, R] = D; + + int paged_B = B; + int paged_K = K; + if (args.mainloop.ptr_page_table != nullptr) { + paged_B = args.mainloop.page_count; + paged_K = args.mainloop.page_size; + } + + auto params_qk_latent = + CollectiveMmaQK::to_underlying_arguments(make_shape(H, K, L, B), + typename CollectiveMmaQK::Arguments{ + args.mainloop.ptr_q_latent, + args.mainloop.stride_q_latent, + args.mainloop.ptr_c_latent, + args.mainloop.stride_c_latent, + }, + nullptr); + + auto params_qk_latent_paged = + CollectiveMmaQK::to_underlying_arguments(make_shape(H, paged_K, L, paged_B), + typename CollectiveMmaQK::Arguments{ + args.mainloop.ptr_q_latent, + args.mainloop.stride_q_latent, + args.mainloop.ptr_c_latent, + args.mainloop.stride_c_latent, + }, + nullptr); + + auto params_qk_rope = + CollectiveMmaQK::to_underlying_arguments(make_shape(H, K, R, B), + typename CollectiveMmaQK::Arguments{ + args.mainloop.ptr_q_rope, + args.mainloop.stride_q_rope, + args.mainloop.ptr_k_rope, + args.mainloop.stride_k_rope, + }, + nullptr); + + auto params_qk_rope_paged = + CollectiveMmaQK::to_underlying_arguments(make_shape(H, paged_K, R, paged_B), + typename CollectiveMmaQK::Arguments{ + args.mainloop.ptr_q_rope, + args.mainloop.stride_q_rope, + args.mainloop.ptr_k_rope, + args.mainloop.stride_k_rope, + }, + nullptr); + + auto stride_c_latent_transpose = select<1, 0, 2>(args.mainloop.stride_c_latent); + auto params_pv_latent = CollectiveMmaPV::to_underlying_arguments( + make_shape(H, L, paged_K, paged_B), + typename CollectiveMmaPV::Arguments{ + args.mainloop.ptr_q_latent, + args.mainloop.stride_q_latent, // dummy, never used + args.mainloop.ptr_c_latent, + stride_c_latent_transpose, + }, + nullptr); + + MainloopParams mainloop_params{params_qk_latent.tma_load_a, params_qk_rope.tma_load_a, + params_qk_latent_paged.tma_load_b, + params_qk_rope_paged.tma_load_b, params_pv_latent.tma_load_b}; + + EpilogueParams epilogue_params; + + epilogue_params.ptr_o = args.epilogue.ptr_o; + epilogue_params.stride_o = args.epilogue.stride_o; + epilogue_params.ptr_lse = args.epilogue.ptr_lse; + epilogue_params.stride_lse = args.epilogue.stride_lse; + epilogue_params.output_scale = args.epilogue.output_scale; + + if (args.split_kv > 1) { + ElementAcc* ptr_o_acc = reinterpret_cast(workspace); + ElementLSE* ptr_lse_acc = + reinterpret_cast(ptr_o_acc + H * L * args.split_kv * B); + epilogue_params.ptr_o_acc = ptr_o_acc; + epilogue_params.ptr_lse_acc = ptr_lse_acc; + + epilogue_params.stride_o_acc = make_tuple(static_cast(0 + L) * args.split_kv, _1{}, + static_cast(0 + H * L) * args.split_kv); + epilogue_params.stride_lse_acc = make_tuple(_1{}, (0 + H) * args.split_kv); + } + + return {args.problem_shape, + args.mainloop, + epilogue_params, + mainloop_params, + TileScheduler::to_underlying_arguments(args.problem_shape, args.hw_info, ClusterShape{}, + args.split_kv), + args.split_kv, + args.ptr_split_kv}; + } + + static size_t get_workspace_size(Arguments const& args) { + ProblemShape problem_shape = args.problem_shape; + auto [H, K, D, B] = problem_shape; + auto [D_latent, D_rope] = D; + auto split_kv = args.split_kv; + return (sizeof(ElementAcc) * D_latent + sizeof(ElementLSE)) * H * split_kv * B; + } + static Status initialize_workspace(Arguments const& /*args*/, void* /*ws*/, + cudaStream_t /*stream*/) { + return Status::kSuccess; + } + + static dim3 get_grid_shape(Params const& params) { + return TileScheduler::get_grid_shape(params.tile_scheduler); + } + + static dim3 get_block_shape() { + dim3 block(MaxThreadsPerBlock, 1, 1); + return block; + } + + static bool can_implement(Arguments const& args) { + if (kIsCpAsync) { + if ((args.mainloop.page_size & (args.mainloop.page_size - 1)) != 0) { + return false; + } + if (args.mainloop.page_size > TileShapeS{}) { + return false; + } + } else { + if (args.mainloop.ptr_page_table != nullptr && args.mainloop.page_size != TileShapeS{}) { + return false; + } + } + if (get<0>(args.problem_shape) != 128) { + return false; + } + if (get<1>(args.problem_shape) <= 0) { + return false; + } + if (args.split_kv <= 0) { + return false; + } + return true; + } + + CUTLASS_DEVICE void operator()(Params const& params, char* smem_raw) { + TileScheduler tile_scheduler(params.tile_scheduler); + + int warp_idx = cutlass::canonical_warp_idx_sync(); + auto role = warp_idx_to_role(warp_idx); + uint32_t lane_predicate = cute::elect_one_sync(); + + uint32_t cta_rank_in_cluster = cute::block_rank_in_cluster(); + int cta_coord_v = cta_rank_in_cluster % size<0>(AtomThrShapeMNK{}); + bool is_mma_leader_cta = cta_coord_v == 0; + + if (role == WarpRole::kLoad && lane_predicate && !kIsCpAsync) { + prefetch_tma_descriptor(params.mainloop_params.tma_load_q_latent.get_tma_descriptor()); + prefetch_tma_descriptor(params.mainloop_params.tma_load_c_latent.get_tma_descriptor()); + prefetch_tma_descriptor(params.mainloop_params.tma_load_q_rope.get_tma_descriptor()); + prefetch_tma_descriptor(params.mainloop_params.tma_load_k_rope.get_tma_descriptor()); + prefetch_tma_descriptor( + params.mainloop_params.tma_load_c_latent_transpose.get_tma_descriptor()); + } + SharedStorage& shared_storage = *reinterpret_cast(smem_raw); + + typename PipelineLoadQK::Params pipeline_load_qk_params; + if (role == WarpRole::kLoad) { + pipeline_load_qk_params.role = PipelineLoadQK::ThreadCategory::Producer; + } + if (role == WarpRole::kMma) { + pipeline_load_qk_params.role = PipelineLoadQK::ThreadCategory::Consumer; + } + if constexpr (kIsCpAsync) { + // we can make our life easier by unconditionally loading blocks + // since we know it'll always be legal + pipeline_load_qk_params.producer_arv_count = + kNumLoadWarps * cutlass::NumThreadsPerWarp * size(AtomThrShapeMNK{}); + } else { + pipeline_load_qk_params.is_leader = + lane_predicate && (role == WarpRole::kLoad) && is_mma_leader_cta; + pipeline_load_qk_params.transaction_bytes = kTransactionsBytesLoadQK; + } + pipeline_load_qk_params.initializing_warp = 0; + PipelineLoadQK pipeline_load_qk(shared_storage.pipelines.load_qk, pipeline_load_qk_params, + ClusterShape{}, /*barrier init*/ cute::true_type{}, + /*mask calc*/ cute::false_type{}); + + typename PipelineS::Params pipeline_mma_s_params; + if (role == WarpRole::kMma) { + pipeline_mma_s_params.role = PipelineS::ThreadCategory::Producer; + } + if (role == WarpRole::kCompute) { + pipeline_mma_s_params.role = PipelineS::ThreadCategory::Consumer; + } + pipeline_mma_s_params.consumer_arv_count = + kNumComputeWarps * cutlass::NumThreadsPerWarp * size(AtomThrShapeMNK{}); + pipeline_mma_s_params.initializing_warp = 1; + PipelineS pipeline_mma_s(shared_storage.pipelines.mma_s, pipeline_mma_s_params, ClusterShape{}, + /*barrier init*/ cute::true_type{}, /*mask calc*/ cute::false_type{}); + + typename PipelineP::Params pipeline_p_mma_params; + if (role == WarpRole::kMma) { + pipeline_p_mma_params.role = PipelineP::ThreadCategory::Consumer; + } + if (role == WarpRole::kCompute) { + pipeline_p_mma_params.role = PipelineP::ThreadCategory::Producer; + } + pipeline_p_mma_params.producer_arv_count = + kNumComputeWarps * cutlass::NumThreadsPerWarp * size(AtomThrShapeMNK{}); + pipeline_p_mma_params.consumer_arv_count = 1; + pipeline_p_mma_params.initializing_warp = 2; + PipelineP pipeline_p_mma(shared_storage.pipelines.p_mma, pipeline_p_mma_params, ClusterShape{}, + /*barrier init*/ cute::true_type{}, /*mask calc*/ cute::false_type{}); + + typename PipelineO::Params pipeline_mma_o_params; + if (role == WarpRole::kMma) { + pipeline_mma_o_params.role = PipelineO::ThreadCategory::Producer; + } + if (role == WarpRole::kCompute) { + pipeline_mma_o_params.role = PipelineO::ThreadCategory::Consumer; + } + pipeline_mma_o_params.consumer_arv_count = + kNumComputeWarps * cutlass::NumThreadsPerWarp * size(AtomThrShapeMNK{}); + pipeline_mma_o_params.initializing_warp = 3; + PipelineO pipeline_mma_o(shared_storage.pipelines.mma_o, pipeline_mma_o_params, ClusterShape{}, + /*barrier init*/ cute::true_type{}, /*mask calc*/ cute::false_type{}); + + typename PipelinePT::Params pipeline_pt_params; + if (role == WarpRole::kLoad) { + pipeline_pt_params.role = PipelinePT::ThreadCategory::Consumer; + } + if (role == WarpRole::kLoadPageTable) { + pipeline_pt_params.role = PipelinePT::ThreadCategory::Producer; + } + pipeline_pt_params.consumer_arv_count = kNumLoadWarps * cutlass::NumThreadsPerWarp; + pipeline_pt_params.producer_arv_count = cutlass::NumThreadsPerWarp; + pipeline_pt_params.initializing_warp = 4; + PipelinePT pipeline_page_table(shared_storage.pipelines.load_page_table, pipeline_pt_params); + + TmemAllocator tmem_allocator; + + pipeline_init_arrive_relaxed(size(ClusterShape{})); + + pipeline_load_qk.init_masks(ClusterShape{}); // do we need an update here for 2Sm? + pipeline_mma_s.init_masks(ClusterShape{}); + pipeline_p_mma.init_masks(ClusterShape{}); + pipeline_mma_o.init_masks(ClusterShape{}); + + typename PipelineLoadQK::PipelineState pipeline_load_qk_consumer_state; + typename PipelineLoadQK::PipelineState pipeline_load_qk_producer_state = + cutlass::make_producer_start_state(); + + typename PipelineS::PipelineState pipeline_mma_s_consumer_state; + typename PipelineS::PipelineState pipeline_mma_s_producer_state = + cutlass::make_producer_start_state(); + + typename PipelineP::PipelineState pipeline_p_mma_consumer_state; + typename PipelineP::PipelineState pipeline_p_mma_producer_state = + cutlass::make_producer_start_state(); + + typename PipelineO::PipelineState pipeline_mma_o_consumer_state; + typename PipelineO::PipelineState pipeline_mma_o_producer_state = + cutlass::make_producer_start_state(); + + typename PipelinePT::PipelineState pipeline_pt_consumer_state; + typename PipelinePT::PipelineState pipeline_pt_producer_state = + cutlass::make_producer_start_state(); + + pipeline_init_wait(size(ClusterShape{})); + + if (role == WarpRole::kLoadPageTable) { + CUTLASS_PRAGMA_NO_UNROLL + for (; tile_scheduler.is_valid(); ++tile_scheduler) { + auto blk_coord = tile_scheduler.get_block_coord(); + auto problem_shape = params.problem_shape; + auto local_split_kv = params.split_kv; + if (params.mainloop.ptr_seq != nullptr) { + get<1>(problem_shape) = params.mainloop.ptr_seq[get<2>(blk_coord)]; + if (params.ptr_split_kv != nullptr) { + local_split_kv = params.ptr_split_kv[get<2>(blk_coord)]; + } + } + if (local_split_kv <= get<3>(blk_coord)) continue; + load_page_table(blk_coord, problem_shape, params.mainloop, shared_storage.tensors, + pipeline_page_table, pipeline_pt_producer_state, local_split_kv); + } + } else if (role == WarpRole::kLoad) { + if constexpr (kIsCpAsync) { + CUTLASS_PRAGMA_NO_UNROLL + for (; tile_scheduler.is_valid(); ++tile_scheduler) { + auto blk_coord = tile_scheduler.get_block_coord(); + auto problem_shape = params.problem_shape; + auto local_split_kv = params.split_kv; + if (params.mainloop.ptr_seq != nullptr) { + get<1>(problem_shape) = params.mainloop.ptr_seq[get<2>(blk_coord)]; + if (params.ptr_split_kv != nullptr) { + local_split_kv = params.ptr_split_kv[get<2>(blk_coord)]; + } + } + if (local_split_kv <= get<3>(blk_coord)) continue; + load_cpasync(blk_coord, problem_shape, params.mainloop, params.mainloop_params, + shared_storage.tensors, pipeline_load_qk, pipeline_load_qk_producer_state, + local_split_kv, + /* must be shared pipe */ + pipeline_page_table, pipeline_pt_consumer_state); + cutlass::arch::NamedBarrier((kNumComputeWarps + kNumLoadWarps) * NumThreadsPerWarp, + kNamedBarrierEpilogue) + .arrive_and_wait(); + } + } else { + if (params.mainloop.ptr_page_table != nullptr) { + CUTLASS_PRAGMA_NO_UNROLL + for (; tile_scheduler.is_valid(); ++tile_scheduler) { + auto blk_coord = tile_scheduler.get_block_coord(); + auto problem_shape = params.problem_shape; + auto local_split_kv = params.split_kv; + if (params.mainloop.ptr_seq != nullptr) { + get<1>(problem_shape) = params.mainloop.ptr_seq[get<2>(blk_coord)]; + if (params.ptr_split_kv != nullptr) { + local_split_kv = params.ptr_split_kv[get<2>(blk_coord)]; + } + } + if (local_split_kv <= get<3>(blk_coord)) continue; + load_tma( + blk_coord, problem_shape, params.mainloop, params.mainloop_params, + shared_storage.tensors, pipeline_load_qk, pipeline_load_qk_producer_state, + pipeline_load_qk, pipeline_load_qk_producer_state, local_split_kv); + cutlass::arch::NamedBarrier((kNumComputeWarps + kNumLoadWarps) * NumThreadsPerWarp, + kNamedBarrierEpilogue) + .arrive_and_wait(); + } + } else { + CUTLASS_PRAGMA_NO_UNROLL + for (; tile_scheduler.is_valid(); ++tile_scheduler) { + auto blk_coord = tile_scheduler.get_block_coord(); + auto problem_shape = params.problem_shape; + auto local_split_kv = params.split_kv; + if (params.mainloop.ptr_seq != nullptr) { + get<1>(problem_shape) = params.mainloop.ptr_seq[get<2>(blk_coord)]; + if (params.ptr_split_kv != nullptr) { + local_split_kv = params.ptr_split_kv[get<2>(blk_coord)]; + } + } + if (local_split_kv <= get<3>(blk_coord)) continue; + load_tma(blk_coord, problem_shape, params.mainloop, params.mainloop_params, + shared_storage.tensors, pipeline_load_qk, + pipeline_load_qk_producer_state, pipeline_load_qk, + pipeline_load_qk_producer_state, local_split_kv); + cutlass::arch::NamedBarrier((kNumComputeWarps + kNumLoadWarps) * NumThreadsPerWarp, + kNamedBarrierEpilogue) + .arrive_and_wait(); + } + } + } + } else if (role == WarpRole::kMma) { + tmem_allocator.allocate(TmemAllocator::Sm100TmemCapacityColumns, + &shared_storage.tmem_base_ptr); + __syncwarp(); + + if (is_mma_leader_cta) { + CUTLASS_PRAGMA_NO_UNROLL + for (; tile_scheduler.is_valid(); ++tile_scheduler) { + auto blk_coord = tile_scheduler.get_block_coord(); + auto problem_shape = params.problem_shape; + auto local_split_kv = params.split_kv; + if (params.mainloop.ptr_seq != nullptr) { + get<1>(problem_shape) = params.mainloop.ptr_seq[get<2>(blk_coord)]; + if (params.ptr_split_kv != nullptr) { + local_split_kv = params.ptr_split_kv[get<2>(blk_coord)]; + } + } + if (local_split_kv <= get<3>(blk_coord)) continue; + mma(blk_coord, problem_shape, shared_storage.tensors, pipeline_load_qk, + pipeline_load_qk_consumer_state, pipeline_load_qk, pipeline_load_qk_consumer_state, + pipeline_mma_s, pipeline_mma_s_producer_state, pipeline_p_mma, + pipeline_p_mma_consumer_state, pipeline_mma_o, pipeline_mma_o_producer_state, + local_split_kv); + } + } + + // cutlass::arch::NamedBarrier((kNumComputeWarps + 1) * NumThreadsPerWarp, + // kNamedBarrierTmemDealloc).arrive_and_wait(); + + // uint32_t free_stage_ptr = shared_storage.tmem_base_ptr; + // tmem_allocator.free(free_stage_ptr, TmemAllocator::Sm100TmemCapacityColumns); + } else if (role == WarpRole::kCompute) { + CUTLASS_PRAGMA_NO_UNROLL + for (; tile_scheduler.is_valid(); ++tile_scheduler) { + auto blk_coord = tile_scheduler.get_block_coord(); + auto problem_shape = params.problem_shape; + auto split_kv = params.split_kv; + auto local_split_kv = split_kv; + if (params.mainloop.ptr_seq != nullptr) { + get<1>(problem_shape) = params.mainloop.ptr_seq[get<2>(blk_coord)]; + if (params.ptr_split_kv != nullptr) { + local_split_kv = params.ptr_split_kv[get<2>(blk_coord)]; + } + } + if (local_split_kv <= get<3>(blk_coord)) continue; + compute(blk_coord, problem_shape, + params.mainloop, // for softmax_scale + params.epilogue, + shared_storage.tensors, // for smem_comm + pipeline_mma_s, pipeline_mma_s_consumer_state, pipeline_p_mma, + pipeline_p_mma_producer_state, pipeline_mma_o, pipeline_mma_o_consumer_state, + local_split_kv); + } + + // cutlass::arch::NamedBarrier((kNumComputeWarps + 1) * NumThreadsPerWarp, + // kNamedBarrierTmemDealloc).arrive(); + } + + cute::cluster_sync(); + cutlass::arch::NamedBarrier((kNumComputeWarps + 1) * NumThreadsPerWarp, + kNamedBarrierTmemDealloc) + .arrive(); + if (role == WarpRole::kMma) { + uint32_t free_stage_ptr = shared_storage.tmem_base_ptr; + tmem_allocator.free(free_stage_ptr, TmemAllocator::Sm100TmemCapacityColumns); + } + } + + template + CUTLASS_DEVICE void load_page_table( + BlkCoord const& blk_coord, ProblemShape const& problem_shape, + MainloopArguments const& mainloop_args, TensorStorage& shared_tensors, + PipelinePT& pipeline_page_table, + typename PipelinePT::PipelineState& pipeline_pt_producer_state, int const& split_kv) { + auto [H, K, D, B] = problem_shape; + int batch_coord = get<2>(blk_coord); + + auto mPT_l = + make_tensor(make_gmem_ptr(mainloop_args.ptr_page_table), + make_shape(mainloop_args.page_count, B), mainloop_args.stride_page_table); + auto mPT = mPT_l(_, batch_coord); + + int k_tile_total = ceil_div(K, TileShapeS{}); + int k_tile_per_cta = ceil_div(k_tile_total, split_kv); + int k_index = get<3>(blk_coord) * k_tile_per_cta; // lower limit + int k_tile_count = max(0, min(k_tile_total, k_index + k_tile_per_cta) - k_index); + if (k_tile_count == 0) { + return; + } + + auto page_size = Pow2{mainloop_args.page_size}; + auto pages_per_tile = Pow2{TileShapeS{} / page_size}; + int thread_idx = threadIdx.x % cutlass::NumThreadsPerWarp; + +#if 1 + for (; k_tile_count > 0; ++k_index, --k_tile_count) { + pipeline_page_table.producer_acquire(pipeline_pt_producer_state); + + // assume a single warp + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < TileShapeS{}; i += cutlass::NumThreadsPerWarp) { + int idx = i + thread_idx; + bool guard = idx < pages_per_tile; + int smem_idx = pipeline_pt_producer_state.index() * TileShapeS::value + idx; + int pt_idx = pages_per_tile * k_index + idx; + + cutlass::arch::cp_async_zfill( + &shared_tensors.smem_page_table[smem_idx], &mPT(pt_idx), guard); + } + + pipeline_page_table.producer_commit(pipeline_pt_producer_state, + cutlass::arch::cpasync_barrier_arrive); + ++pipeline_pt_producer_state; + } +#endif + } + + struct Gather { + int& page_table_stage; + Pow2 pages_per_tile; + const int* __restrict__ smem_page_table; + + CUTLASS_DEVICE int operator()(int idx) const { + return smem_page_table[page_table_stage * TileShapeS::value + idx % pages_per_tile]; + } + + CUTLASS_DEVICE friend void print(Gather const&) { printf(""); } + }; + + template + CUTLASS_DEVICE void load_cpasync( + BlkCoord const& blk_coord, ProblemShape const& problem_shape, + MainloopArguments const& mainloop_args, MainloopParams const& mainloop_params, + TensorStorage& shared_tensors, PipelineLoadQK& pipeline_load, + typename PipelineLoadQK::PipelineState& pipeline_load_producer_state, int const& split_kv, + PipelinePT& pipeline_page_table, + typename PipelinePT::PipelineState& pipeline_pt_consumer_state) { + auto [H, K, D, B] = problem_shape; + auto [D_latent, D_rope] = D; + + using X = Underscore; + + int k_tile_total = ceil_div(K, TileShapeS{}); + int k_tile_per_cta = ceil_div(k_tile_total, split_kv); + int k_index = get<3>(blk_coord) * k_tile_per_cta; // lower limit + int k_tile_count = max(0, min(k_tile_total, k_index + k_tile_per_cta) - k_index); + if (k_tile_count == 0) { + return; + } + + // partition all tensors + auto mQL = make_tensor(make_gmem_ptr(mainloop_args.ptr_q_latent), make_shape(H, D_latent, B), + mainloop_args.stride_q_latent); + auto mQR = make_tensor(make_gmem_ptr(mainloop_args.ptr_q_rope), make_shape(H, D_rope, B), + mainloop_args.stride_q_rope); + + int paged_B = mainloop_args.page_count; + auto paged_K = Pow2{mainloop_args.page_size}; + auto mPT_l = make_tensor(make_gmem_ptr(mainloop_args.ptr_page_table), make_shape(paged_B, B), + mainloop_args.stride_page_table); + + int batch_coord = get<2>(blk_coord); + auto mPT = mPT_l(_, batch_coord); + + auto gQL = local_tile(mQL, TileShapeQK{}, make_coord(_, _, _), Step<_1, X, _1>{}); + auto gQR = local_tile(mQR, TileShapeQK{}, make_coord(_, _, _), Step<_1, X, _1>{}); + + ThrMMA cta_mma_qk = TiledMmaQK{}.get_slice(get<0>(blk_coord) % size(AtomThrShapeMNK{})); + ThrMMA cta_mma_pv = TiledMmaPV{}.get_slice(get<0>(blk_coord) % size(AtomThrShapeMNK{})); + + auto tSgQL = cta_mma_qk.partition_A(gQL); + auto tSgQR = cta_mma_qk.partition_A(gQR); + + Tensor sQ = make_tensor(make_smem_ptr(shared_tensors.smem_q.begin()), SmemLayoutQ{}); + Tensor sKC = make_tensor(make_smem_ptr(shared_tensors.smem_kc.begin()), SmemLayoutKC{}); + Tensor sVC = make_tensor(make_smem_ptr(shared_tensors.smem_vc.begin()), SmemLayoutVC{}); + + auto make_copy_for = [](auto sT) { + auto rT_a = sT.layout()(_, _, _, _0{}); + auto rT = make_ordered_layout(shape(rT_a), stride(rT_a)); + auto threads = Int{}; + auto values = Int{}; + return make_cotiled_copy( + Copy_Atom, Element>{}, + make_ordered_layout(make_shape(threads, values), make_stride(_1{}, _0{})), rT); + }; + + // like cute::copy, but makes sure we do all page table lookups first + auto copy_split = [](auto atom, auto src, auto dst) { + auto src_v = group_modes<1, rank_v>(src); + auto dst_v = group_modes<1, rank_v>(dst); + + auto src_v_ptrs = make_tensor(size<1>(src_v)); + for (int i = 0; i < size<1>(src_v); i++) { + src_v_ptrs(i) = &src_v(_0{}, i); + } + + for (int i = 0; i < size<1>(src_v); i++) { + auto src_v_i = make_tensor(make_gmem_ptr(src_v_ptrs(i)), make_shape(shape<0>(src_v)), + make_stride(make_stride(_1{}, _0{}))); + atom.call(src_v_i, dst_v(_, i)); + } + }; + + auto tiled_copy_q = make_copy_for(sQ); + auto tiled_copy_kc = make_copy_for(sKC); + auto tiled_copy_vc = make_copy_for(sVC); + + auto thr_copy_q = + tiled_copy_q.get_thread_slice(threadIdx.x % (kNumLoadWarps * cutlass::NumThreadsPerWarp)); + auto thr_copy_kc = + tiled_copy_kc.get_thread_slice(threadIdx.x % (kNumLoadWarps * cutlass::NumThreadsPerWarp)); + auto thr_copy_vc = + tiled_copy_vc.get_thread_slice(threadIdx.x % (kNumLoadWarps * cutlass::NumThreadsPerWarp)); + + auto tQsQ = thr_copy_q.partition_D(sQ); + auto tQgQL = thr_copy_q.partition_S(tSgQL); + auto tQgQR = thr_copy_q.partition_S(tSgQR); + + auto tKCsKC = thr_copy_kc.partition_D(sKC); + auto tVCsVC = thr_copy_vc.partition_D(sVC); + + auto pipeline_pt_release_state = pipeline_pt_consumer_state; + + int page_table_stage = -1; + Pow2 pages_per_tile{TileShapeS{} / paged_K}; + const int* __restrict__ smem_page_table = shared_tensors.smem_page_table.begin(); + Gather gather{page_table_stage, pages_per_tile, smem_page_table}; + + auto mCL = make_tensor( + make_gmem_ptr(mainloop_args.ptr_c_latent), + ComposedLayout{ + make_layout(make_shape(make_shape(paged_K, paged_B), _1{}), + make_stride(make_stride(get<0>(mainloop_args.stride_c_latent), + example::CustomStride( + gather, get<2>(mainloop_args.stride_c_latent))), + get<1>(mainloop_args.stride_c_latent))), + make_coord(_0{}, _0{}), make_identity_layout(make_shape(paged_K * paged_B, D_latent))}); + + auto mKR = make_tensor( + make_gmem_ptr(mainloop_args.ptr_k_rope), + ComposedLayout{ + make_layout(make_shape(make_shape(paged_K, paged_B), _1{}), + make_stride(make_stride(get<0>(mainloop_args.stride_k_rope), + example::CustomStride( + gather, get<2>(mainloop_args.stride_k_rope))), + get<1>(mainloop_args.stride_k_rope))), + make_coord(_0{}, _0{}), make_identity_layout(make_shape(paged_K * paged_B, D_latent))}); + + auto mCLT = make_tensor( + make_gmem_ptr(mainloop_args.ptr_c_latent), + ComposedLayout{ + make_layout( + make_shape(_1{}, make_shape(paged_K, paged_B)), + make_stride(get<1>(mainloop_args.stride_c_latent), + make_stride(get<0>(mainloop_args.stride_c_latent), + example::CustomStride( + gather, get<2>(mainloop_args.stride_c_latent))))), + make_coord(_0{}, _0{}), make_identity_layout(make_shape(D_latent, paged_K * paged_B))}); + + auto gCL = local_tile(mCL, TileShapeQK{}, make_coord(_, _, _), Step{}); + auto gKR = local_tile(mKR, TileShapeQK{}, make_coord(_, _, _), Step{}); + auto gCLT = local_tile(mCLT, TileShapePV{}, make_coord(_, _, _), Step{}); + + auto tSgCL = cta_mma_qk.partition_B(gCL); + auto tSgKR = cta_mma_qk.partition_B(gKR); + auto tOgCLT = cta_mma_pv.partition_B(gCLT); + + auto tKCgCL = thr_copy_kc.partition_S(tSgCL); + auto tKCgKR = thr_copy_kc.partition_S(tSgKR); + auto tVCgCLT = thr_copy_vc.partition_S(tOgCLT); + + // latent is first in memory, so let's load it first always + // startup: alternate Q and K, set tx count appropriately, for k_idx = 0 + auto& pipeline_acquire_state = pipeline_load_producer_state; + auto pipeline_commit_state = pipeline_acquire_state; + int pipeline_offset = 0; + + for (int i = 0; i < StagesPV; i++) { + cutlass::arch::cp_async_fence(); + } + + auto load_stage = [&](auto fn) { + pipeline_load.producer_acquire(pipeline_acquire_state); + fn(pipeline_acquire_state.index()); + cutlass::arch::cp_async_fence(); + + ++pipeline_acquire_state; + ++pipeline_offset; + + if (pipeline_offset == StagesPV - 1) { + cutlass::arch::cp_async_wait(); + pipeline_load.producer_commit(pipeline_commit_state); + ++pipeline_commit_state; + --pipeline_offset; + } + }; + + pipeline_page_table.consumer_wait(pipeline_pt_consumer_state); + page_table_stage = pipeline_pt_consumer_state.index(); + ++pipeline_pt_consumer_state; + + // each Q/K tile consists of rope and latent + for (int i = 0; i < IterationsQKLatent; i++) { + load_stage([&](int index) { + cute::copy(tiled_copy_q, tQgQL(_, _, _, _, _0{}, i, batch_coord), tQsQ(_, _, _, _, i)); + copy_split(tiled_copy_kc, tKCgCL(_, _, _, _, k_index, i), tKCsKC(_, _, _, _, index)); + }); + } + + for (int i = 0; i < IterationsQKRope; i++) { + load_stage([&](int index) { + cute::copy(tiled_copy_q, tQgQR(_, _, _, _, _0{}, i, batch_coord), + tQsQ(_, _, _, _, IterationsQKLatent + i)); + copy_split(tiled_copy_kc, tKCgKR(_, _, _, _, k_index, i), tKCsKC(_, _, _, _, index)); + }); + } + + k_index += 1; + k_tile_count -= 1; + + // assume k_tile_count >= 1 + // perform K+Q load here + CUTLASS_PRAGMA_NO_UNROLL + while (k_tile_count > 0) { + pipeline_page_table.consumer_wait(pipeline_pt_consumer_state); + page_table_stage = pipeline_pt_consumer_state.index(); + ++pipeline_pt_consumer_state; + + for (int i = 0; i < IterationsQKLatent; i++) { + load_stage([&](int index) { + copy_split(tiled_copy_kc, tKCgCL(_, _, _, _, k_index, i), tKCsKC(_, _, _, _, index)); + }); + } + + for (int i = 0; i < IterationsQKRope; i++) { + load_stage([&](int index) { + copy_split(tiled_copy_kc, tKCgKR(_, _, _, _, k_index, i), tKCsKC(_, _, _, _, index)); + }); + } + + page_table_stage = pipeline_pt_release_state.index(); + + for (int i = 0; i < IterationsPV_K; i++) { + for (int j = 0; j < IterationsPV_N; j++) { + load_stage([&](int index) { + copy_split(tiled_copy_vc, tVCgCLT(_, _, _, _, j, IterationsPV_K * (k_index - 1) + i), + tVCsVC(_, _, _, _, index)); + }); + } + } + + pipeline_page_table.consumer_release(pipeline_pt_release_state); + ++pipeline_pt_release_state; + + k_index += 1; + k_tile_count -= 1; + } + + page_table_stage = pipeline_pt_release_state.index(); + + for (int i = 0; i < IterationsPV_K; i++) { + for (int j = 0; j < IterationsPV_N; j++) { + load_stage([&](int index) { + copy_split(tiled_copy_vc, tVCgCLT(_, _, _, _, j, IterationsPV_K * (k_index - 1) + i), + tVCsVC(_, _, _, _, index)); + }); + } + } + + pipeline_page_table.consumer_release(pipeline_pt_release_state); + ++pipeline_pt_release_state; + + while (pipeline_offset > 0) { + cutlass::arch::cp_async_fence(); + + cutlass::arch::cp_async_wait(); + pipeline_load.producer_commit(pipeline_commit_state); + ++pipeline_commit_state; + --pipeline_offset; + } + + cutlass::arch::cp_async_wait<0>(); + } + + template + CUTLASS_DEVICE void load_tma( + BlkCoord const& blk_coord, ProblemShape const& problem_shape, + MainloopArguments const& mainloop_args, MainloopParams const& mainloop_params, + TensorStorage& shared_tensors, PipelineLoadQK& pipeline_load_qk, + typename PipelineLoadQK::PipelineState& pipeline_load_qk_producer_state, + PipelineLoadPV& pipeline_load_pv, + typename PipelineLoadPV::PipelineState& pipeline_load_pv_producer_state, + int const& split_kv) { + auto [H, K, D, B] = problem_shape; + auto [D_latent, D_rope] = D; + + int k_tile_total = ceil_div(K, TileShapeS{}); + int k_tile_per_cta = ceil_div(k_tile_total, split_kv); + int k_index = get<3>(blk_coord) * k_tile_per_cta; // lower limit + int k_tile_count = max(0, min(k_tile_total, k_index + k_tile_per_cta) - k_index); + if (k_tile_count == 0) { + return; + } + + using X = Underscore; + + // partition all tensors + auto mQL = mainloop_params.tma_load_q_latent.get_tma_tensor(make_shape(H, D_latent, B)); + auto mQR = mainloop_params.tma_load_q_rope.get_tma_tensor(make_shape(H, D_rope, B)); + + int paged_B = B; + int paged_K = K; + if constexpr (kIsPaged) { + paged_B = mainloop_args.page_count; + paged_K = mainloop_args.page_size; + } + auto mPT_l = make_tensor(make_gmem_ptr(mainloop_args.ptr_page_table), make_shape(paged_B, B), + mainloop_args.stride_page_table); + + auto mCL = + mainloop_params.tma_load_c_latent.get_tma_tensor(make_shape(paged_K, D_latent, paged_B)); + auto mKR = mainloop_params.tma_load_k_rope.get_tma_tensor(make_shape(paged_K, D_rope, paged_B)); + + auto mCLT = mainloop_params.tma_load_c_latent_transpose.get_tma_tensor( + make_shape(D_latent, paged_K, paged_B)); + + auto gQL = local_tile(mQL, TileShapeQK{}, make_coord(_, _, _), Step<_1, X, _1>{}); + auto gQR = local_tile(mQR, TileShapeQK{}, make_coord(_, _, _), Step<_1, X, _1>{}); + + auto gCL = local_tile(mCL, TileShapeQK{}, make_coord(_, _, _), Step{}); + auto gKR = local_tile(mKR, TileShapeQK{}, make_coord(_, _, _), Step{}); + auto gCLT = local_tile(mCLT, TileShapePV{}, make_coord(_, _, _), Step{}); + + ThrMMA cta_mma_qk = TiledMmaQK{}.get_slice(get<0>(blk_coord) % size(AtomThrShapeMNK{})); + ThrMMA cta_mma_pv = TiledMmaPV{}.get_slice(get<0>(blk_coord) % size(AtomThrShapeMNK{})); + + auto tSgQL = cta_mma_qk.partition_A(gQL); + auto tSgQR = cta_mma_qk.partition_A(gQR); + + auto tSgCL = cta_mma_qk.partition_B(gCL); + auto tSgKR = cta_mma_qk.partition_B(gKR); + + auto tOgCLT = cta_mma_pv.partition_B(gCLT); + + Tensor sQ = make_tensor(make_smem_ptr(shared_tensors.smem_q.begin()), SmemLayoutQ{}); + Tensor sKC = make_tensor(make_smem_ptr(shared_tensors.smem_kc.begin()), SmemLayoutKC{}); + Tensor sVC = make_tensor(make_smem_ptr(shared_tensors.smem_vc.begin()), SmemLayoutVC{}); + + auto [tQLgQL_mkl, tQsQ] = + tma_partition(mainloop_params.tma_load_q_latent, _0{}, make_layout(_1{}), + group_modes<0, 3>(sQ), group_modes<0, 3>(tSgQL)); + + auto [tQRgQR_mkl, tQsQ_ignore] = + tma_partition(mainloop_params.tma_load_q_rope, _0{}, make_layout(_1{}), + group_modes<0, 3>(sQ), group_modes<0, 3>(tSgQR)); + + auto [tCLgCL_nkl, tKCsKC] = + tma_partition(mainloop_params.tma_load_c_latent, _0{}, make_layout(_1{}), + group_modes<0, 3>(sKC), group_modes<0, 3>(tSgCL)); + + auto [tKRgKR_nkl, tKCsKC_ignore] = + tma_partition(mainloop_params.tma_load_k_rope, _0{}, make_layout(_1{}), + group_modes<0, 3>(sKC), group_modes<0, 3>(tSgKR)); + + auto [tCLTgCLT_nkl, tVCsVC] = + tma_partition(mainloop_params.tma_load_c_latent_transpose, _0{}, make_layout(_1{}), + group_modes<0, 3>(sVC), group_modes<0, 3>(tOgCLT)); + + uint16_t mcast_mask = 0; + + int batch_coord = get<2>(blk_coord); + Tensor tQLgQL = tQLgQL_mkl(_, _, _, batch_coord); + Tensor tQRgQR = tQRgQR_mkl(_, _, _, batch_coord); + + auto mPT = mPT_l(_, batch_coord); + + Tensor tCLgCL = tCLgCL_nkl(_, _, _, _); + Tensor tKRgKR = tKRgKR_nkl(_, _, _, _); + + // careful: stage and k are swapped here! + Tensor tCLTgCLT = tCLTgCLT_nkl(_, _, _, _); + + // latent is first in memory, so let's load it first always + // startup: alternate Q and K, set tx count appropriately, for k_idx = 0 + + // each Q/K tile consists of rope and latent + for (int i = 0; i < IterationsQKLatent; i++) { + pipeline_load_qk.producer_expect_transaction(pipeline_load_qk_producer_state, + kTransactionsBytesLoadExtraQ); + pipeline_load_qk.producer_acquire(pipeline_load_qk_producer_state); + auto tma_barrier = pipeline_load_qk.producer_get_barrier(pipeline_load_qk_producer_state); + + if (cute::elect_one_sync()) { + // expect the extra bytes + // load_qk ql + cute::copy(mainloop_params.tma_load_q_latent.with(*tma_barrier, mcast_mask), + tQLgQL(_, _0{}, i), tQsQ(_, i)); + // load_qk cl + if constexpr (kIsPaged) { + cute::copy(mainloop_params.tma_load_c_latent.with(*tma_barrier, mcast_mask), + tCLgCL(_, _0{}, i, mPT(k_index)), + tKCsKC(_, pipeline_load_qk_producer_state.index())); + } else { + cute::copy(mainloop_params.tma_load_c_latent.with(*tma_barrier, mcast_mask), + tCLgCL(_, k_index, i, batch_coord), + tKCsKC(_, pipeline_load_qk_producer_state.index())); + } + } + ++pipeline_load_qk_producer_state; + } + + for (int i = 0; i < IterationsQKRope; i++) { + pipeline_load_qk.producer_expect_transaction(pipeline_load_qk_producer_state, + kTransactionsBytesLoadExtraQ); + pipeline_load_qk.producer_acquire(pipeline_load_qk_producer_state); + auto tma_barrier = pipeline_load_qk.producer_get_barrier(pipeline_load_qk_producer_state); + + if (cute::elect_one_sync()) { + // expect the extra bytes + // load_qk ql + cute::copy(mainloop_params.tma_load_q_rope.with(*tma_barrier, mcast_mask), + tQRgQR(_, _0{}, i), tQsQ(_, i + IterationsQKLatent)); + // load_qk cl + if constexpr (kIsPaged) { + cute::copy(mainloop_params.tma_load_k_rope.with(*tma_barrier, mcast_mask), + tKRgKR(_, _0{}, i, mPT(k_index)), + tKCsKC(_, pipeline_load_qk_producer_state.index())); + } else { + cute::copy(mainloop_params.tma_load_k_rope.with(*tma_barrier, mcast_mask), + tKRgKR(_, k_index, i, batch_coord), + tKCsKC(_, pipeline_load_qk_producer_state.index())); + } + } + ++pipeline_load_qk_producer_state; + } + + k_index += 1; + k_tile_count -= 1; + + // assume k_tile_count >= 1 + // perform K+Q load here + CUTLASS_PRAGMA_NO_UNROLL + while (k_tile_count > 0) { + // perform K load + for (int i = 0; i < IterationsQKLatent; i++) { + pipeline_load_qk.producer_acquire(pipeline_load_qk_producer_state); + auto tma_barrier = pipeline_load_qk.producer_get_barrier(pipeline_load_qk_producer_state); + + if (cute::elect_one_sync()) { + // load_qk cl + if constexpr (kIsPaged) { + cute::copy(mainloop_params.tma_load_c_latent.with(*tma_barrier, mcast_mask), + tCLgCL(_, _0{}, i, mPT(k_index)), + tKCsKC(_, pipeline_load_qk_producer_state.index())); + } else { + cute::copy(mainloop_params.tma_load_c_latent.with(*tma_barrier, mcast_mask), + tCLgCL(_, k_index, i, batch_coord), + tKCsKC(_, pipeline_load_qk_producer_state.index())); + } + } + ++pipeline_load_qk_producer_state; + } + + for (int i = 0; i < IterationsQKRope; i++) { + pipeline_load_qk.producer_acquire(pipeline_load_qk_producer_state); + auto tma_barrier = pipeline_load_qk.producer_get_barrier(pipeline_load_qk_producer_state); + + if (cute::elect_one_sync()) { + // load_qk cl + if constexpr (kIsPaged) { + cute::copy(mainloop_params.tma_load_k_rope.with(*tma_barrier, mcast_mask), + tKRgKR(_, _0{}, i, mPT(k_index)), + tKCsKC(_, pipeline_load_qk_producer_state.index())); + } else { + cute::copy(mainloop_params.tma_load_k_rope.with(*tma_barrier, mcast_mask), + tKRgKR(_, k_index, i, batch_coord), + tKCsKC(_, pipeline_load_qk_producer_state.index())); + } + } + ++pipeline_load_qk_producer_state; + } + + // prefetch next K load to keep busy while we transpose-load from cache + const int kPrefetchDistance = 1; + for (int i = 0; i < IterationsQKLatent; i++) { + if (cute::elect_one_sync()) { + if constexpr (kIsPaged) { + if (k_tile_count > kPrefetchDistance) { + cute::prefetch(mainloop_params.tma_load_c_latent, + tCLgCL(_, _0{}, i, mPT(k_index + kPrefetchDistance))); + } + } else { + cute::prefetch(mainloop_params.tma_load_c_latent, + tCLgCL(_, k_index + kPrefetchDistance, i, batch_coord)); + } + } + } + + for (int i = 0; i < IterationsQKRope; i++) { + if (cute::elect_one_sync()) { + if constexpr (kIsPaged) { + if (k_tile_count > kPrefetchDistance) { + cute::prefetch(mainloop_params.tma_load_k_rope, + tKRgKR(_, _0{}, i, mPT(k_index + kPrefetchDistance))); + } + } else { + cute::prefetch(mainloop_params.tma_load_k_rope, + tKRgKR(_, k_index + kPrefetchDistance, i, batch_coord)); + } + } + } + + // perform V load (k_idx - 1) + + for (int i = 0; i < IterationsPV_K; i++) { + for (int j = 0; j < IterationsPV_N; j++) { + pipeline_load_pv.producer_acquire(pipeline_load_pv_producer_state); + auto tma_barrier = pipeline_load_pv.producer_get_barrier(pipeline_load_pv_producer_state); + + if (cute::elect_one_sync()) { + // load_pv cl + // note the transpose in indices! + // note we are off-by-one on k_index + if constexpr (kIsPaged) { + cute::copy(mainloop_params.tma_load_c_latent_transpose.with( + *tma_barrier, mcast_mask, cute::TMA::CacheHintSm100::EVICT_FIRST), + tCLTgCLT(_, j, i, mPT(k_index - 1)), + tVCsVC(_, pipeline_load_pv_producer_state.index())); + } else { + cute::copy(mainloop_params.tma_load_c_latent_transpose.with( + *tma_barrier, mcast_mask, cute::TMA::CacheHintSm100::EVICT_FIRST), + tCLTgCLT(_, j, IterationsPV_K * (k_index - 1) + i, batch_coord), + tVCsVC(_, pipeline_load_pv_producer_state.index())); + } + } + ++pipeline_load_pv_producer_state; + } + } + + k_index += 1; + k_tile_count -= 1; + } + + for (int i = 0; i < IterationsPV_K; i++) { + for (int j = 0; j < IterationsPV_N; j++) { + pipeline_load_pv.producer_acquire(pipeline_load_pv_producer_state); + auto tma_barrier = pipeline_load_pv.producer_get_barrier(pipeline_load_pv_producer_state); + + if (cute::elect_one_sync()) { + // load_pv cl + // note the transpose in indices + // note we are off-by-one on k_index + + if constexpr (kIsPaged) { + cute::copy(mainloop_params.tma_load_c_latent_transpose.with( + *tma_barrier, mcast_mask, cute::TMA::CacheHintSm100::EVICT_FIRST), + tCLTgCLT(_, j, i, mPT(k_index - 1)), + tVCsVC(_, pipeline_load_pv_producer_state.index())); + } else { + cute::copy(mainloop_params.tma_load_c_latent_transpose.with( + *tma_barrier, mcast_mask, cute::TMA::CacheHintSm100::EVICT_FIRST), + tCLTgCLT(_, j, IterationsPV_K * (k_index - 1) + i, batch_coord), + tVCsVC(_, pipeline_load_pv_producer_state.index())); + } + } + ++pipeline_load_pv_producer_state; + } + } + } + + template + CUTLASS_DEVICE void mma(BlkCoord const& blk_coord, ProblemShape const& problem_shape, + TensorStorage& shared_tensors, PipelineLoadQK& pipeline_load_qk, + typename PipelineLoadQK::PipelineState& pipeline_load_qk_consumer_state, + PipelineLoadPV& pipeline_load_pv, + typename PipelineLoadPV::PipelineState& pipeline_load_pv_consumer_state, + PipelineS& pipeline_mma_s, + typename PipelineS::PipelineState& pipeline_mma_s_producer_state, + PipelineP& pipeline_p_mma, + typename PipelineP::PipelineState& pipeline_p_mma_consumer_state, + PipelineO& pipeline_mma_o, + typename PipelineO::PipelineState& pipeline_mma_o_producer_state, + int const& split_kv) { + auto [H, K, D, B] = problem_shape; + + int k_tile_total = ceil_div(K, TileShapeS{}); + int k_tile_per_cta = ceil_div(k_tile_total, split_kv); + int k_index = get<3>(blk_coord) * k_tile_per_cta; // lower limit + int k_tile_count = max(0, min(k_tile_total, k_index + k_tile_per_cta) - k_index); + if (k_tile_count == 0) { + return; + } + + // mma init + Tensor sQ = make_tensor(make_smem_ptr(shared_tensors.smem_q.begin()), SmemLayoutQ{}); + Tensor sKC = make_tensor(make_smem_ptr(shared_tensors.smem_kc.begin()), SmemLayoutKC{}); + Tensor sVC = make_tensor(make_smem_ptr(shared_tensors.smem_vc.begin()), SmemLayoutVC{}); + Tensor sP = make_tensor(make_smem_ptr((Element*)shared_tensors.smem_p.begin()), SmemLayoutP{}); + + Tensor tSrQ = TiledMmaQK::make_fragment_A(sQ); + Tensor tSrKC = TiledMmaQK::make_fragment_B(sKC); + Tensor tOrP = TiledMmaPV::make_fragment_A(sP); + Tensor tOrVC = TiledMmaPV::make_fragment_B(sVC); + + TiledMmaQK tiled_mma_qk; + TiledMmaPV tiled_mma_pv; + + Tensor tStS = partition_fragment_C(tiled_mma_qk, select<0, 1>(TileShapeQK{})); + Tensor tOtO = partition_fragment_C(tiled_mma_pv, select<0, 1>(TileShapePV{})); + + tiled_mma_pv.accumulate_ = UMMA::ScaleOut::Zero; + + pipeline_mma_s.producer_acquire(pipeline_mma_s_producer_state); + + // Mma S0 S1 O0 S2 O1 ... Sn On-1 On + // S0 ownership -- ----- -- -- + // S1 ownership -- ----- ---- + // O ownership -- -- ---- -- + + tiled_mma_qk.accumulate_ = UMMA::ScaleOut::Zero; + for (int i = 0; i < IterationsQK; i++) { + pipeline_load_qk.consumer_wait(pipeline_load_qk_consumer_state); + int read_stage = pipeline_load_qk_consumer_state.index(); + + tStS.data() = uint32_t(pipeline_mma_s_producer_state.index() == 0 ? TmemAllocation::kS0 + : TmemAllocation::kS1); + + CUTLASS_PRAGMA_UNROLL + for (int k_block = 0; k_block < size<2>(tSrQ); ++k_block) { + cute::gemm(tiled_mma_qk, tSrQ(_, _, k_block, i), tSrKC(_, _, k_block, read_stage), tStS); + tiled_mma_qk.accumulate_ = UMMA::ScaleOut::One; + } + + pipeline_load_qk.consumer_release(pipeline_load_qk_consumer_state); + ++pipeline_load_qk_consumer_state; + } + + pipeline_mma_s.producer_commit(pipeline_mma_s_producer_state); + ++pipeline_mma_s_producer_state; + + k_tile_count -= 1; + + CUTLASS_PRAGMA_NO_UNROLL + while (k_tile_count > 0) { + pipeline_mma_s.producer_acquire(pipeline_mma_s_producer_state); + tiled_mma_qk.accumulate_ = UMMA::ScaleOut::Zero; + for (int i = 0; i < IterationsQK; i++) { + pipeline_load_qk.consumer_wait(pipeline_load_qk_consumer_state); + int read_stage = pipeline_load_qk_consumer_state.index(); + + tStS.data() = uint32_t(pipeline_mma_s_producer_state.index() == 0 ? TmemAllocation::kS0 + : TmemAllocation::kS1); + + CUTLASS_PRAGMA_UNROLL + for (int k_block = 0; k_block < size<2>(tSrQ); ++k_block) { + cute::gemm(tiled_mma_qk, tSrQ(_, _, k_block, i), tSrKC(_, _, k_block, read_stage), tStS); + tiled_mma_qk.accumulate_ = UMMA::ScaleOut::One; + } + + pipeline_load_qk.consumer_release(pipeline_load_qk_consumer_state); + ++pipeline_load_qk_consumer_state; + } + + pipeline_mma_s.producer_commit(pipeline_mma_s_producer_state); + ++pipeline_mma_s_producer_state; + + pipeline_mma_o.producer_acquire(pipeline_mma_o_producer_state); + pipeline_p_mma.consumer_wait(pipeline_p_mma_consumer_state); + + for (int i = 0; i < IterationsPV_K; i++) { + auto acc_flag = tiled_mma_pv.accumulate_; + for (int j = 0; j < IterationsPV_N; j++) { + pipeline_load_pv.consumer_wait(pipeline_load_pv_consumer_state); + + int read_stage = pipeline_load_pv_consumer_state.index(); + + tOtO.data() = uint32_t(TmemAllocation::kO0) + j * uint32_t(TmemAllocation::kSizeAccO); + tiled_mma_pv.accumulate_ = acc_flag; + + CUTLASS_PRAGMA_UNROLL + for (int k_block = 0; k_block < size<2>(tOrP); ++k_block) { + cute::gemm(tiled_mma_pv, + tOrP(_, _, k_block, make_coord(i, pipeline_p_mma_consumer_state.index())), + tOrVC(_, _, k_block, read_stage), tOtO); + tiled_mma_pv.accumulate_ = UMMA::ScaleOut::One; + } + + pipeline_load_pv.consumer_release(pipeline_load_pv_consumer_state); + ++pipeline_load_pv_consumer_state; + } + } + + pipeline_p_mma.consumer_release(pipeline_p_mma_consumer_state); + ++pipeline_p_mma_consumer_state; + pipeline_mma_o.producer_commit(pipeline_mma_o_producer_state); + ++pipeline_mma_o_producer_state; + + --k_tile_count; + } + + pipeline_mma_o.producer_acquire(pipeline_mma_o_producer_state); + pipeline_p_mma.consumer_wait(pipeline_p_mma_consumer_state); + + for (int i = 0; i < IterationsPV_K; i++) { + auto acc_flag = tiled_mma_pv.accumulate_; + for (int j = 0; j < IterationsPV_N; j++) { + pipeline_load_pv.consumer_wait(pipeline_load_pv_consumer_state); + + int read_stage = pipeline_load_pv_consumer_state.index(); + + tOtO.data() = uint32_t(TmemAllocation::kO0) + j * uint32_t(TmemAllocation::kSizeAccO); + tiled_mma_pv.accumulate_ = acc_flag; + + CUTLASS_PRAGMA_UNROLL + for (int k_block = 0; k_block < size<2>(tOrP); ++k_block) { + cute::gemm(tiled_mma_pv, + tOrP(_, _, k_block, make_coord(i, pipeline_p_mma_consumer_state.index())), + tOrVC(_, _, k_block, read_stage), tOtO); + tiled_mma_pv.accumulate_ = UMMA::ScaleOut::One; + } + + pipeline_load_pv.consumer_release(pipeline_load_pv_consumer_state); + ++pipeline_load_pv_consumer_state; + } + } + + pipeline_p_mma.consumer_release(pipeline_p_mma_consumer_state); + ++pipeline_p_mma_consumer_state; + pipeline_mma_o.producer_commit(pipeline_mma_o_producer_state); + ++pipeline_mma_o_producer_state; + } + + template + CUTLASS_DEVICE void softmax(IsLastTile const& is_last_tile, ElementAcc& row_max, + ElementAcc& row_sum, ElementAcc& correction_factor, + ProblemShape const& problem_shape, + MainloopArguments const& mainloop_args, TensorStorage& shared_tensors, + int k_index, uint32_t tmem_s, int smem_p_index) { + auto load_op = cute::SM100_TMEM_LOAD_32dp32b32x{}; + + TiledMmaQK tiled_mma_qk; + + Tensor tStS = partition_fragment_C(tiled_mma_qk, select<0, 1>(TileShapeQK{})); + tStS.data() = tmem_s; + + CUTE_STATIC_ASSERT_V(shape<1>(tStS) == _1{}); + CUTE_STATIC_ASSERT_V(shape<2>(tStS) == _1{}); + Tensor tAcc = tStS(make_coord(_, _), _0{}, _0{}); + + Tensor cS = make_identity_tensor(take<0, 2>(CtaShapeQK{})); + + auto tiled_t2r = make_tmem_copy(load_op, tAcc); + auto thread_idx = threadIdx.x % size(tiled_t2r); + + auto thread_t2r = tiled_t2r.get_slice(thread_idx); + Tensor tTR_cS = thread_t2r.partition_D(cS); + Tensor tTR_rAcc = make_tensor(shape(tTR_cS)); + + Tensor tTR_rS_frag = make_tensor(shape(tTR_rAcc)); + const int AlignmentS = 4; + Tensor tTR_tAcc = thread_t2r.partition_S(tAcc); + Tensor tTR_rAcc_vec = recast>(tTR_rAcc); + Tensor tTR_rS_vec = recast>(tTR_rS_frag); + + // load s + copy(tiled_t2r, tTR_tAcc, tTR_rAcc); + + if (is_last_tile) { + for (int i = 0; i < size(tTR_rAcc); i++) { + if (get<1>(tTR_cS(i)) + TileShapeS{} * k_index >= get<1>(problem_shape)) { + tTR_rAcc(i) = -std::numeric_limits::infinity(); + } + } + } + + // max + ElementAcc row_max_new = row_max; + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(tTR_rAcc); i += 1) { + row_max_new = ::fmax(row_max_new, tTR_rAcc(i)); + } + + // for 2x2 dp, reduce here + if constexpr (kWarpsInN > 1) { + shared_tensors.smem_exchange[threadIdx.x] = row_max_new; + cutlass::arch::NamedBarrier(kNumComputeWarps * NumThreadsPerWarp, kNamedBarrierExchange) + .sync(); + // (64, 2) shape + int peer_index = (threadIdx.x + 64) % 128; + row_max_new = cutlass::max(row_max_new, shared_tensors.smem_exchange[peer_index]); + } + +#ifndef B2B + // find correction factor + ElementAcc softmax_scale_log2 = mainloop_args.softmax_scale * static_cast(M_LOG2E); + correction_factor = ::exp2f(softmax_scale_log2 * (row_max - row_max_new)); + row_max = row_max_new; + + // softmax + ElementAcc row_max_scale_log2 = row_max * softmax_scale_log2; + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(tTR_rAcc); i++) { + tTR_rAcc(i) = ::exp2f(softmax_scale_log2 * tTR_rAcc(i) - row_max_scale_log2); + } +#endif + + // quantize + cutlass::NumericArrayConverter epilogue_op; + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(tTR_rAcc_vec); i++) { + tTR_rS_vec(i) = epilogue_op(tTR_rAcc_vec(i)); + } + + Tensor sP = make_tensor(make_smem_ptr((Element*)shared_tensors.smem_p.begin()), SmemLayoutP{})( + _, _, _, make_coord(_, smem_p_index)); + + Tensor tOcP = TiledMmaPV{}.get_slice(_0{}).partition_A(cS); + + // have a mapping for each thread to coord + // find identical mapping to coords for the MMA + auto l = make_ordered_layout( + make_shape(make_shape(_64{}, _2{}), make_shape(_16{}, TileShapeS{} / _32{})), + make_stride(make_stride(_0{}, _3{}), make_stride(_1{}, _2{}))); + auto sP_ = as_position_independent_swizzle_tensor(sP); + copy_aligned(tTR_rS_frag, sP_.compose(l)(threadIdx.x, _)); + + // sum + row_sum *= correction_factor; + + static_assert(cute::is_same_v); + auto tTR_rAcc_float2 = recast(tTR_rAcc); + auto sums = make_tensor(_4{}); + static_assert(size(tTR_rAcc_float2) % size(sums) == 0); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(sums); i++) { + sums(i) = tTR_rAcc_float2(i); + } + CUTLASS_PRAGMA_UNROLL + for (int i = size(sums); i < size(tTR_rAcc_float2); i += size(sums)) { + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < size(sums); j++) { + cute::add(sums(j), sums(j), tTR_rAcc_float2(i + j)); + } + } + CUTLASS_PRAGMA_UNROLL + for (int i = 1; i < size(sums); i *= 2) { + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < size(sums); j += 2 * i) { + cute::add(sums(j), sums(j), sums(j + i)); + } + } + row_sum += sums(0).x + sums(0).y; + } + + CUTLASS_DEVICE void rescale(ElementAcc correction_factor, uint32_t tmem_o) { + // for b2b gemm, do nothing +#ifndef B2B + auto load_op = cute::SM100_TMEM_LOAD_32dp32b32x{}; + auto store_op = TMEM::tmem_load_to_store(load_op); + + TiledMmaPV tiled_mma_pv; + + Tensor tOtO = partition_fragment_C(tiled_mma_pv, select<0, 1>(TileShapePV{})); + tOtO.data() = tmem_o; + + CUTE_STATIC_ASSERT_V(shape<1>(tOtO) == _1{}); + CUTE_STATIC_ASSERT_V(shape<2>(tOtO) == _1{}); + Tensor tAcc = tOtO(make_coord(_, _), _0{}, _0{}); + + auto cta_tiler_pv = take<0, 2>(typename CollectiveMmaPV::CtaShape_MNK{}); + Tensor gO = make_tensor(make_gmem_ptr((ElementAcc*)nullptr), cta_tiler_pv, make_stride(0, 0)); + + auto tiled_t2r = make_tmem_copy(load_op, tAcc); + auto tiled_r2t = make_tmem_copy(store_op, tAcc); + auto thread_idx = threadIdx.x % size(tiled_t2r); + + auto thread_t2r = tiled_t2r.get_slice(thread_idx); + auto thread_r2t = tiled_r2t.get_slice(thread_idx); + Tensor tTR_gO = thread_t2r.partition_D(gO); + Tensor tTR_rAcc = make_tensor(shape(tTR_gO)); + + Tensor tTR_tAcc = thread_t2r.partition_S(tAcc); + + // load o + copy(tiled_t2r, tTR_tAcc, tTR_rAcc); + + // multiply by correction factor + float2 correction_factor_vec = make_float2(correction_factor, correction_factor); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(tTR_rAcc); i += 2) { + float2 in = make_float2(tTR_rAcc(i + 0), tTR_rAcc(i + 1)); + float2 out; + cute::mul(out, in, correction_factor_vec); + tTR_rAcc(i + 0) = out.x; + tTR_rAcc(i + 1) = out.y; + } + + // store o + copy(tiled_r2t, tTR_rAcc, tTR_tAcc); +#endif + } + + template + CUTLASS_DEVICE void epilogue(ElementAcc& row_max, ElementAcc& row_sum, BlkCoord const& cta_coord, + ProblemShape const& problem_shape, + MainloopArguments const& mainloop_args, + EpilogueParams const& epilogue_args, TensorStorage& shared_tensors, + uint32_t tmem_o, int const& split_kv) { + auto load_op = cute::SM100_TMEM_LOAD_32dp32b32x{}; + + TiledMmaPV tiled_mma_pv; + + Tensor tOtO = + TiledMmaPV::make_fragment_C(partition_shape_C(TiledMmaPV{}, take<0, 2>(TileShapePV{}))); + tOtO.data() = tmem_o; + + CUTE_STATIC_ASSERT_V(shape<1>(tOtO) == _1{}); + CUTE_STATIC_ASSERT_V(shape<2>(tOtO) == _1{}); + Tensor tAcc = tOtO(make_coord(_, _), _0{}, _0{}); + + auto [H, K, D, B] = problem_shape; + auto [D_latent, D_rope] = D; + if (epilogue_args.ptr_o_acc != nullptr) { + using ElementOutAcc = ElementAcc; + constexpr auto AlignmentOutAcc = 128 / cute::sizeof_bits_v; + Tensor mO = make_tensor(make_gmem_ptr(epilogue_args.ptr_o_acc + get<3>(cta_coord) * D_latent), + make_shape(H, D_latent, B), epilogue_args.stride_o_acc); + auto cta_tiler_pv = take<0, 2>(typename CollectiveMmaPV::CtaShape_MNK{}); + Tensor gO = local_tile(mO, cta_tiler_pv, take<0, 3>(cta_coord)); + + auto tiled_t2r = make_tmem_copy(load_op, tAcc); + auto thread_idx = threadIdx.x % size(tiled_t2r); + + auto thread_t2r = tiled_t2r.get_slice(thread_idx); + Tensor tTR_gO = thread_t2r.partition_D(gO); + Tensor tTR_rAcc = make_tensor(shape(tTR_gO)); + + Tensor tTR_rO_frag = make_tensor(shape(tTR_rAcc)); + Tensor tTR_rO_src = recast>(coalesce(tTR_rO_frag)); + Tensor tR2G_rO_dst = recast>(coalesce(tTR_gO)); + Tensor tTR_tAcc = thread_t2r.partition_S(tAcc); + + copy(tiled_t2r, tTR_tAcc, tTR_rAcc); + + cutlass::epilogue::thread::LinearCombination< + ElementOutAcc, 1, ElementAcc, ElementAcc, + cutlass::epilogue::thread::ScaleType::OnlyAlphaScaling> + epilogue_op({epilogue_args.output_scale / row_sum}); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(tTR_rAcc); i++) { + tTR_rO_frag(i) = epilogue_op(tTR_rAcc(i)); + } + + copy(tTR_rO_src, tR2G_rO_dst); + +#ifndef B2B + + // compute LSE + ElementAcc lse = cutlass::fast_log(row_sum) + mainloop_args.softmax_scale * row_max; + + // store LSE + Tensor mLSE = make_tensor(make_gmem_ptr(epilogue_args.ptr_lse_acc + H * get<3>(cta_coord)), + make_shape(H, B), epilogue_args.stride_lse_acc); + Tensor gLSE = local_tile(mLSE, append<3>(cta_tiler_pv, _1{}), take<0, 3>(cta_coord), + Step<_1, Underscore, _1>{}); + // for 2x2 dp, this must be conditional and the index is wrong + if (!kIs2Sm || (threadIdx.x < 64)) { + gLSE(threadIdx.x) = lse; + } +#endif + } else { + Tensor mO = make_tensor(make_gmem_ptr(epilogue_args.ptr_o), make_shape(H, D_latent, B), + epilogue_args.stride_o); + auto cta_tiler_pv = take<0, 2>(typename CollectiveMmaPV::CtaShape_MNK{}); + Tensor gO = local_tile(mO, cta_tiler_pv, take<0, 3>(cta_coord)); + + auto tiled_t2r = make_tmem_copy(load_op, tAcc); + auto thread_idx = threadIdx.x % size(tiled_t2r); + + auto thread_t2r = tiled_t2r.get_slice(thread_idx); + Tensor tTR_gO = thread_t2r.partition_D(gO); + Tensor tTR_rAcc = make_tensor(shape(tTR_gO)); + + Tensor tTR_rO_frag = make_tensor(shape(tTR_rAcc)); + Tensor tTR_rO_src = recast>(coalesce(tTR_rO_frag)); + Tensor tR2G_rO_dst = recast>(coalesce(tTR_gO)); + Tensor tTR_tAcc = thread_t2r.partition_S(tAcc); + + copy(tiled_t2r, tTR_tAcc, tTR_rAcc); + + cutlass::epilogue::thread::LinearCombination< + ElementOut, 1, ElementAcc, ElementAcc, + cutlass::epilogue::thread::ScaleType::OnlyAlphaScaling> + epilogue_op({epilogue_args.output_scale / row_sum}); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(tTR_rAcc); i++) { + tTR_rO_frag(i) = epilogue_op(tTR_rAcc(i)); + } + + copy(tTR_rO_src, tR2G_rO_dst); + +#ifndef B2B + if (epilogue_args.ptr_lse != nullptr) { + // compute LSE + ElementAcc lse = cutlass::fast_log(row_sum) + mainloop_args.softmax_scale * row_max; + + // store LSE + Tensor mLSE = make_tensor(make_gmem_ptr(epilogue_args.ptr_lse), make_shape(H, B), + epilogue_args.stride_lse); + Tensor gLSE = local_tile(mLSE, append<3>(cta_tiler_pv, _1{}), take<0, 3>(cta_coord), + Step<_1, Underscore, _1>{}); + + // for 2x2 dp, this must be conditional and the index is wrong + if (!kIs2Sm || (threadIdx.x < 64)) { + gLSE(threadIdx.x) = lse; + } + } +#endif + } + } + + template + CUTLASS_DEVICE void compute( + CtaCoord const& cta_coord, ProblemShape const& problem_shape, + MainloopArguments const& mainloop_args, EpilogueParams const& epilogue_args, + TensorStorage& shared_tensors, PipelineS& pipeline_mma_s, + typename PipelineS::PipelineState& pipeline_mma_s_consumer_state, PipelineP& pipeline_p_mma, + typename PipelineP::PipelineState& pipeline_p_mma_producer_state, PipelineO& pipeline_mma_o, + typename PipelineO::PipelineState& pipeline_mma_o_consumer_state, int const& split_kv) { + auto [H, K, D, B] = problem_shape; + + int k_tile_total = ceil_div(K, TileShapeS{}); + int k_tile_per_cta = ceil_div(k_tile_total, split_kv); + int k_index = get<3>(cta_coord) * k_tile_per_cta; // lower limit + int k_tile_count = max(0, min(k_tile_total, k_index + k_tile_per_cta) - k_index); + if (k_tile_count == 0) { + // if we return early, we have to make sure we release the load warp + cutlass::arch::NamedBarrier((kNumComputeWarps + kNumLoadWarps) * NumThreadsPerWarp, + kNamedBarrierEpilogue) + .arrive(); + + return; + } + int k_index_final = k_tile_total - 1; + + ElementAcc row_max = -std::numeric_limits::infinity(); + ElementAcc row_sum = 0; + ElementAcc correction_factor = 1; + + pipeline_p_mma.producer_acquire(pipeline_p_mma_producer_state); + pipeline_mma_s.consumer_wait(pipeline_mma_s_consumer_state); + + auto dispatch_bool = [](bool b, auto fn) { + if (b) { + fn(cute::true_type{}); + } else { + fn(cute::false_type{}); + } + }; + + // softmax s0 -> p0 + dispatch_bool(k_index == k_index_final, [&](auto is_last_tile) { + softmax(is_last_tile, row_max, row_sum, correction_factor, problem_shape, mainloop_args, + shared_tensors, k_index, + uint32_t(pipeline_mma_s_consumer_state.index() == 0 ? TmemAllocation::kS0 + : TmemAllocation::kS1), + pipeline_p_mma_producer_state.index()); + }); + + k_index += 1; + + cutlass::arch::fence_view_async_tmem_load(); + cutlass::arch::fence_view_async_shared(); + pipeline_mma_s.consumer_release(pipeline_mma_s_consumer_state); + ++pipeline_mma_s_consumer_state; + pipeline_p_mma.producer_commit(pipeline_p_mma_producer_state); + ++pipeline_p_mma_producer_state; + + k_tile_count -= 1; + + CUTLASS_PRAGMA_NO_UNROLL + while (k_tile_count > 0) { + pipeline_p_mma.producer_acquire(pipeline_p_mma_producer_state); + pipeline_mma_s.consumer_wait(pipeline_mma_s_consumer_state); + + // softmax s1 -> p1 + dispatch_bool(k_index == k_index_final, [&](auto is_last_tile) { + softmax(is_last_tile, row_max, row_sum, correction_factor, problem_shape, mainloop_args, + shared_tensors, k_index, + uint32_t(pipeline_mma_s_consumer_state.index() == 0 ? TmemAllocation::kS0 + : TmemAllocation::kS1), + pipeline_p_mma_producer_state.index()); + }); + + cutlass::arch::fence_view_async_tmem_load(); + cutlass::arch::fence_view_async_shared(); + pipeline_mma_s.consumer_release(pipeline_mma_s_consumer_state); + ++pipeline_mma_s_consumer_state; + pipeline_p_mma.producer_commit(pipeline_p_mma_producer_state); + ++pipeline_p_mma_producer_state; + + pipeline_mma_o.consumer_wait(pipeline_mma_o_consumer_state); + + // rescale + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < IterationsPV_N; j++) { + rescale(correction_factor, + uint32_t(TmemAllocation::kO0) + j * uint32_t(TmemAllocation::kSizeAccO)); + } + + cutlass::arch::fence_view_async_tmem_store(); + pipeline_mma_o.consumer_release(pipeline_mma_o_consumer_state); + ++pipeline_mma_o_consumer_state; + + --k_tile_count; + k_index += 1; + } + + pipeline_mma_o.consumer_wait(pipeline_mma_o_consumer_state); + +#ifdef B2B + row_sum = 1; +#else + if constexpr (kWarpsInN > 1) { + // reduce row_sum if needed (for 2x2 dp) + shared_tensors.smem_exchange[threadIdx.x] = row_sum; + cutlass::arch::NamedBarrier(kNumComputeWarps * NumThreadsPerWarp, kNamedBarrierExchange) + .sync(); + // (64, 2) shape + int peer_index = (threadIdx.x + 64) % 128; + row_sum += shared_tensors.smem_exchange[peer_index]; + } +#endif + + cutlass::arch::NamedBarrier((kNumComputeWarps + kNumLoadWarps) * NumThreadsPerWarp, + kNamedBarrierEpilogue) + .arrive(); + + // epilogue + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < IterationsPV_N; j++) { + epilogue(row_max, row_sum, replace<1>(cta_coord, j), problem_shape, mainloop_args, + epilogue_args, shared_tensors, + uint32_t(TmemAllocation::kO0) + j * uint32_t(TmemAllocation::kSizeAccO), split_kv); + } + + cutlass::arch::fence_view_async_tmem_load(); + pipeline_mma_o.consumer_release(pipeline_mma_o_consumer_state); + ++pipeline_mma_o_consumer_state; + } +}; + +/////////////////////////////////////////////////////////////////////////////// + +} // namespace cutlass::fmha::kernel diff --git a/include/flashinfer/attention/blackwell/kernel/sm100_mla_tile_scheduler.hpp b/include/flashinfer/attention/blackwell/kernel/sm100_mla_tile_scheduler.hpp new file mode 100644 index 0000000000000000000000000000000000000000..f096c6d537a9dd724f2c0cbd82de990b804f25ba --- /dev/null +++ b/include/flashinfer/attention/blackwell/kernel/sm100_mla_tile_scheduler.hpp @@ -0,0 +1,151 @@ +/*************************************************************************************************** + * Copyright (c) 2024 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +#pragma once + +#include "cutlass/cutlass.h" +#include "cutlass/fast_math.h" +#include "cutlass/kernel_hardware_info.h" + +namespace cutlass::fmha::kernel { + +//////////////////////////////////////////////////////////////////////////////// + +struct Sm100MlaIndividualTileScheduler { + struct Params { + dim3 grid; + }; + + bool valid_ = true; + + CUTLASS_DEVICE + Sm100MlaIndividualTileScheduler(Params const&) {} + + template + static Params to_underlying_arguments(ProblemShape const& problem_shape, + KernelHardwareInfo hw_info, + ClusterShape const& cluster_shape, int const& split_kv) { + using namespace cute; + dim3 grid(get<0>(cluster_shape), get<3>(problem_shape) /* Batch */, + split_kv /*Maximum Split KV*/); + return Params{grid}; + } + + static dim3 get_grid_shape(Params const& params) { return params.grid; } + + CUTLASS_DEVICE + bool is_valid() { return valid_; } + + CUTLASS_DEVICE + auto get_block_coord() { + using namespace cute; + return make_coord(blockIdx.x, _0{}, blockIdx.y, blockIdx.z); + } + + CUTLASS_DEVICE + Sm100MlaIndividualTileScheduler& operator++() { + valid_ = false; + return *this; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +struct Sm100MlaPersistentTileScheduler { + struct Params { + int num_blocks; + FastDivmod divmod_m_block; + FastDivmod divmod_b; + FastDivmod divmod_split_kv; + KernelHardwareInfo hw_info; + }; + + int block_idx = 0; + Params params; + + CUTLASS_DEVICE + Sm100MlaPersistentTileScheduler(Params const& params) : block_idx(blockIdx.x), params(params) {} + + template + static Params to_underlying_arguments(ProblemShape const& problem_shape, + KernelHardwareInfo hw_info, + ClusterShape const& cluster_shape, int const& split_kv) { + using namespace cute; + // Get SM count if needed, otherwise use user supplied SM count + int sm_count = hw_info.sm_count; + if (sm_count <= 1 || sm_count % size<0>(cluster_shape) != 0) { + CUTLASS_TRACE_HOST( + " WARNING: Arguments do not include a valid SM count.\n" + " For optimal performance, populate the arguments KernelHardwareInfo struct with the SM " + "count."); + sm_count = KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id); + } + + CUTLASS_TRACE_HOST("to_underlying_arguments(): Setting persistent grid SM count to " + << sm_count); + hw_info.sm_count = sm_count; + + int num_m_blocks = size<0>(cluster_shape); + int num_blocks = num_m_blocks * get<3>(problem_shape) /* Batch */; + num_blocks *= split_kv; /* Maximum Split KV*/ + + return Params{num_blocks, {num_m_blocks}, {get<3>(problem_shape)}, {split_kv}, hw_info}; + } + + static dim3 get_grid_shape(Params const& params) { + dim3 grid(std::min(params.num_blocks, params.hw_info.sm_count), 1, 1); + return grid; + } + + CUTLASS_DEVICE + bool is_valid() { return block_idx < params.num_blocks; } + + CUTLASS_DEVICE + auto get_block_coord() { + using namespace cute; + int block_decode = block_idx; + int m_block, bidb, n_split_kv; + params.divmod_m_block(block_decode, m_block, block_decode); + params.divmod_b(block_decode, bidb, block_decode); + params.divmod_split_kv(block_decode, n_split_kv, block_decode); + return make_coord(m_block, _0{}, bidb, n_split_kv); + } + + CUTLASS_DEVICE + Sm100MlaPersistentTileScheduler& operator++() { + block_idx += gridDim.x; + return *this; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace cutlass::fmha::kernel diff --git a/include/flashinfer/attention/blackwell/plan.cuh b/include/flashinfer/attention/blackwell/plan.cuh new file mode 100644 index 0000000000000000000000000000000000000000..fe44c4bd175c2dd6922be630a6846f3bd7164553 --- /dev/null +++ b/include/flashinfer/attention/blackwell/plan.cuh @@ -0,0 +1,175 @@ +/* + * Copyright (c) 2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include + +#include "../../utils.cuh" + +namespace flashinfer { + +union alignas(8) CostIndex { + struct { + int bucket_idx; + float cost; + }; + long long packed; +}; + +__device__ __forceinline__ CostIndex min(CostIndex a, CostIndex b) { + return a.cost < b.cost || (a.cost == b.cost && a.bucket_idx < b.bucket_idx) ? a : b; +} + +__device__ __forceinline__ CostIndex get_min_cost_index(CostIndex* warp_min_cost, + CostIndex cost_index, int num_buckets) { +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + CostIndex other; + other.packed = __shfl_xor_sync(0xffffffff, cost_index.packed, offset); + cost_index = min(cost_index, other); + } + if (threadIdx.x % 32 == 0) { + warp_min_cost[threadIdx.x / 32] = cost_index; + } + __syncthreads(); + if (threadIdx.x < 32) { + cost_index = threadIdx.x * 32 < num_buckets + ? warp_min_cost[threadIdx.x] + : CostIndex{threadIdx.x * 32, cuda::std::numeric_limits::infinity()}; +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + CostIndex other; + other.packed = __shfl_xor_sync(0xffffffff, cost_index.packed, offset); + cost_index = min(cost_index, other); + } + if (threadIdx.x == 0) { + warp_min_cost[0] = cost_index; + } + } + __syncthreads(); + return warp_min_cost[0]; +} + +__global__ void plan_kernel(int* qo_segment_offsets, int* kv_segment_offsets, int* qo_lens, + int* kv_lens, int* work_indptr, int* qo_tile_indices, int* head_indices, + int* batch_indices, int qo_tile_size, int batch_size, int num_heads, + int num_buckets, bool causal) { + __shared__ CostIndex warp_min_cost[32]; + constexpr int MAX_BUCKET_SIZE = 256; + using BlockScan = cub::BlockScan; + __shared__ typename BlockScan::TempStorage temp_storage; + // first round, calculate the work count for each bucket + CostIndex thread_local_cost_index = {threadIdx.x, 0.f}; + int thread_local_work_counter = 0; + if (threadIdx.x >= num_buckets) { + thread_local_cost_index.cost = cuda::std::numeric_limits::infinity(); + } + + for (int head_idx = 0; head_idx < num_heads; ++head_idx) { + for (int batch_idx = 0; batch_idx < batch_size; ++batch_idx) { + int qo_len = qo_lens ? qo_lens[batch_idx] + : qo_segment_offsets[batch_idx + 1] - qo_segment_offsets[batch_idx]; + int kv_len = kv_lens ? kv_lens[batch_idx] + : kv_segment_offsets[batch_idx + 1] - kv_segment_offsets[batch_idx]; + int num_qo_tiles = ceil_div(qo_len, qo_tile_size); + for (int qo_tile_idx = num_qo_tiles - 1; qo_tile_idx >= 0; --qo_tile_idx) { + auto min_cost_index = + get_min_cost_index(warp_min_cost, thread_local_cost_index, num_buckets); + int bucket_idx = min_cost_index.bucket_idx; + float cost = min_cost_index.cost; + if (bucket_idx == threadIdx.x) { + thread_local_cost_index.cost += + causal ? kv_len - (num_qo_tiles - qo_tile_idx - 1) * qo_tile_size : kv_len; + thread_local_work_counter++; + } + } + } + } + __syncthreads(); + // compute exclusive prefix sum of + int thread_local_work_indptr = 0; + BlockScan(temp_storage).ExclusiveSum(thread_local_work_counter, thread_local_work_indptr); + __syncthreads(); + if (threadIdx.x < num_buckets) { + work_indptr[threadIdx.x] = thread_local_work_indptr; + } + if (threadIdx.x + 1 == num_buckets) { + work_indptr[num_buckets] = thread_local_work_indptr + thread_local_work_counter; + } + + // second round, write qo_tile_idx, head_idx, batch_idx to the output + thread_local_work_counter = 0; + if (threadIdx.x >= num_buckets) { + thread_local_cost_index.cost = cuda::std::numeric_limits::infinity(); + } else { + thread_local_cost_index.cost = 0.f; + } + for (int head_idx = 0; head_idx < num_heads; ++head_idx) { + for (int batch_idx = 0; batch_idx < batch_size; ++batch_idx) { + int qo_len = qo_lens ? qo_lens[batch_idx] + : qo_segment_offsets[batch_idx + 1] - qo_segment_offsets[batch_idx]; + int kv_len = kv_lens ? kv_lens[batch_idx] + : kv_segment_offsets[batch_idx + 1] - kv_segment_offsets[batch_idx]; + int num_qo_tiles = ceil_div(qo_len, qo_tile_size); + for (int qo_tile_idx = num_qo_tiles - 1; qo_tile_idx >= 0; --qo_tile_idx) { + auto min_cost_index = + get_min_cost_index(warp_min_cost, thread_local_cost_index, num_buckets); + int bucket_idx = min_cost_index.bucket_idx; + float cost = min_cost_index.cost; + if (bucket_idx == threadIdx.x) { + thread_local_cost_index.cost += + causal ? kv_len - (num_qo_tiles - qo_tile_idx - 1) * qo_tile_size : kv_len; + qo_tile_indices[thread_local_work_indptr + thread_local_work_counter] = qo_tile_idx; + head_indices[thread_local_work_indptr + thread_local_work_counter] = head_idx; + batch_indices[thread_local_work_indptr + thread_local_work_counter] = batch_idx; + thread_local_work_counter++; + } + } + } + } +#if (__CUDACC_VER_MAJOR__ >= 12 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.launch_dependents;"); +#endif +} + +cudaError_t plan_kernel_wrapper(int* qo_segment_offsets, int* kv_segment_offsets, int* qo_lens, + int* kv_lens, int* work_indptr, int* qo_tile_indices, + int* head_indices, int* batch_indices, int qo_tile_size, + int batch_size, int num_heads, int num_buckets, bool causal, + bool enable_pdl, cudaStream_t stream) { + if (enable_pdl) { + cudaLaunchConfig_t config; + config.gridDim = 1; + config.blockDim = 256; + config.dynamicSmemBytes = 0; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = true; + config.numAttrs = 1; + config.attrs = attrs; + FLASHINFER_CUDA_CALL( + cudaLaunchKernelEx(&config, plan_kernel, qo_segment_offsets, kv_segment_offsets, qo_lens, + kv_lens, work_indptr, qo_tile_indices, head_indices, batch_indices, + qo_tile_size, batch_size, num_heads, num_buckets, causal)); + } else { + plan_kernel<<<1, 256, 0, stream>>>(qo_segment_offsets, kv_segment_offsets, qo_lens, kv_lens, + work_indptr, qo_tile_indices, head_indices, batch_indices, + qo_tile_size, batch_size, num_heads, num_buckets, causal); + FLASHINFER_CUDA_CALL(cudaGetLastError()); + } + return cudaSuccess; +} + +} // namespace flashinfer diff --git a/include/flashinfer/attention/cascade.cuh b/include/flashinfer/attention/cascade.cuh new file mode 100644 index 0000000000000000000000000000000000000000..8fb5e6b9192dc3027c0c10c7321d2368aa3abf3b --- /dev/null +++ b/include/flashinfer/attention/cascade.cuh @@ -0,0 +1,786 @@ +/*! + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef FLASHINFER_CASCADE_CUH_ +#define FLASHINFER_CASCADE_CUH_ + +#include "../cp_async.cuh" +#include "../math.cuh" +#include "../utils.cuh" +#include "state.cuh" + +namespace flashinfer { + +using cp_async::PrefetchMode; +using cp_async::SharedMemFillMode; + +/*! + * \brief The CUDA kernel that merges the self-attention state of two index sets A and B. + * \tparam vec_size The vector size used in the kernel. + * \tparam DTypeIn The data type of v_a and v_b. + * \tparam DTypeO The data type of v_merged. + * \param v_a The partial v of index set A. (n, h, d) + * \param s_a The logsumexp value of index set A. (n, h) + * \param v_b The partial v of index set B. (n, h, d) + * \param s_b The logsumexp value of index set B. (n, h) + * \param v_merged The merged v of index set A union B. (n, h, d) + * \param s_merged The merged logsumexp value of index set A union B. (n, h) + * \param num_heads The number of heads of v_a and v_b. + * \param head_dim The dimension of each head. + * \note Both s_a and s_b are logsumexp values with base 2. + */ +template +__global__ void MergeStateKernel(DTypeIn* __restrict__ v_a, float* __restrict__ s_a, + DTypeIn* __restrict__ v_b, float* __restrict__ s_b, + DTypeO* __restrict__ v_merged, float* __restrict__ s_merged, + uint32_t num_heads, uint32_t head_dim) { + uint32_t tx = threadIdx.x, ty = threadIdx.y; + uint32_t pos = blockIdx.x; + uint32_t head_idx = ty; + + float s_a_val = s_a[pos * num_heads + head_idx]; + float s_b_val = s_b[pos * num_heads + head_idx]; + float s_max = max(s_a_val, s_b_val); + s_a_val = math::ptx_exp2(s_a_val - s_max); + s_b_val = math::ptx_exp2(s_b_val - s_max); + float a_scale = s_a_val / (s_a_val + s_b_val); + float b_scale = s_b_val / (s_a_val + s_b_val); + vec_t v_a_vec, v_b_vec, v_merged_vec; + v_a_vec.cast_load(v_a + (pos * num_heads + head_idx) * head_dim + tx * vec_size); + v_b_vec.cast_load(v_b + (pos * num_heads + head_idx) * head_dim + tx * vec_size); +#pragma unroll + for (uint32_t i = 0; i < vec_size; ++i) { + v_merged_vec[i] = a_scale * v_a_vec[i] + b_scale * v_b_vec[i]; + } + v_merged_vec.cast_store(v_merged + (pos * num_heads + head_idx) * head_dim + tx * vec_size); + if (s_merged != nullptr) { + s_merged[pos * num_heads + head_idx] = math::ptx_log2(s_a_val + s_b_val) + s_max; + } +} + +/*! + * \brief The CUDA kernel that merges the self-attention state with another state in-place. + * \tparam vec_size The vector size used in the kernel. + * \tparam DType The data type of v and v_other. + * \param v The partial v to be updated in-place. (n, h, d) + * \param s The logsumexp value to be updated in-place. (n, h) + * \param v_other The other v to be merged. (n, h, d) + * \param s_other The other logsumexp value to be merged. (n, h) + * \param mask Optional mask of whether to merge given sequences or not. (n) + * \param num_heads The number of heads of v and v_other. + * \param head_dim The dimension of each head. + * \note Both s and s_other are logsumexp values with base 2. + */ +template +__global__ void MergeStateInPlaceKernel(DType* __restrict__ v, float* __restrict__ s, + DType* __restrict__ v_other, float* __restrict__ s_other, + uint8_t* __restrict__ mask, uint32_t num_heads, + uint32_t head_dim) { + uint32_t pos = blockIdx.x; + + if (mask != nullptr && mask[pos] == 0) return; + + uint32_t tx = threadIdx.x, ty = threadIdx.y; + uint32_t head_idx = ty; + + float s_val = s[pos * num_heads + head_idx]; + float s_other_val = s_other[pos * num_heads + head_idx]; + float s_max = max(s_val, s_other_val); + s_val = math::ptx_exp2(s_val - s_max); + s_other_val = math::ptx_exp2(s_other_val - s_max); + float scale = s_val / (s_val + s_other_val); + float other_scale = s_other_val / (s_val + s_other_val); + vec_t v_vec, v_other_vec; + v_vec.cast_load(v + (pos * num_heads + head_idx) * head_dim + tx * vec_size); + v_other_vec.cast_load(v_other + (pos * num_heads + head_idx) * head_dim + tx * vec_size); +#pragma unroll + for (uint32_t i = 0; i < vec_size; ++i) { + v_vec[i] = scale * v_vec[i] + other_scale * v_other_vec[i]; + } + v_vec.cast_store(v + (pos * num_heads + head_idx) * head_dim + tx * vec_size); + if (s != nullptr) { + s[pos * num_heads + head_idx] = math::ptx_log2(s_val + s_other_val) + s_max; + } +} + +template +__device__ __forceinline__ void threadblock_sync_state(state_t& st, DTypeIn* v_smem, + float* s_smem, + const uint32_t tx = threadIdx.x, + const uint32_t ty = threadIdx.y) { + constexpr uint32_t head_dim = vec_size * bdx; + st.o.cast_store(v_smem + ty * head_dim + tx * vec_size); + s_smem[ty] = st.get_lse(); + st.init(); + __syncthreads(); + +#pragma unroll + for (uint32_t iter = 0; iter < bdy; ++iter) { + float s = s_smem[iter]; + vec_t v; + v.cast_load(v_smem + iter * head_dim + tx * vec_size); + st.merge(v, s, 1); + } +} + +template +__device__ __forceinline__ void warp_sync_state(state_t& st, DTypeIn* v_smem, + float* s_smem, const uint32_t tx = threadIdx.x, + const uint32_t ty = threadIdx.y) { + constexpr uint32_t head_dim = vec_size * bdx; + st.o.cast_store(v_smem + ty * head_dim + tx * vec_size); + s_smem[ty] = st.get_lse(); + st.init(); + __syncwarp(); + +#pragma unroll + for (uint32_t iter = 0; iter < bdy; ++iter) { + float s = s_smem[iter]; + vec_t v; + v.cast_load(v_smem + iter * head_dim + tx * vec_size); + st.merge(v, s, 1); + } +} + +template +__device__ __forceinline__ void threadblock_sum(vec_t& v, DTypeIn* v_smem) { + const uint32_t tx = threadIdx.x, ty = threadIdx.y; + constexpr uint32_t head_dim = vec_size * bdx; + v.cast_store(v_smem + ty * head_dim + tx * vec_size); + v.fill(DTypeIn(0.f)); + __syncthreads(); + +#pragma unroll + for (uint32_t iter = 0; iter < bdy; ++iter) { + vec_t v_iter; + v_iter.cast_load(v_smem + iter * head_dim + tx * vec_size); +#pragma unroll + for (uint32_t i = 0; i < vec_size; ++i) { + v[i] += v_iter[i]; + } + } +} + +template +__global__ void AttentionSumKernel(DTypeIn* __restrict__ V, DTypeO* __restrict__ v_sum, + uint32_t num_index_sets, uint32_t num_heads, uint32_t head_dim) { + uint32_t tx = threadIdx.x, ty = threadIdx.y; + uint32_t pos = blockIdx.x; + uint32_t head_idx = ty; + + if (num_index_sets == 0) { + vec_t v; + v.fill(DTypeO(0.f)); + v.store(v_sum + (pos * num_heads + head_idx) * head_dim + tx * vec_size); + return; + } + + if (num_index_sets == 1) { + vec_t v; + v.cast_load(V + (pos * num_heads + head_idx) * head_dim + tx * vec_size); + v.store(v_sum + (pos * num_heads + head_idx) * head_dim + tx * vec_size); + return; + } + + vec_t v_sum_vec; + v_sum_vec.fill(0.f); +#pragma unroll 2 + for (uint32_t iter = 0; iter < num_index_sets; ++iter) { + vec_t v; + v.cast_load(V + ((pos * num_index_sets + iter) * num_heads + head_idx) * head_dim + + tx * vec_size); +#pragma unroll + for (uint32_t i = 0; i < vec_size; ++i) { + v_sum_vec[i] += v[i]; + } + } + + v_sum_vec.cast_store(v_sum + (pos * num_heads + head_idx) * head_dim + tx * vec_size); +} + +template +__global__ void MergeStatesKernel(DTypeIn* __restrict__ V, float* __restrict__ S, + DTypeO* __restrict__ v_merged, float* __restrict__ s_merged, + uint32_t num_index_sets, uint32_t num_heads, uint32_t head_dim) { + uint32_t tx = threadIdx.x, ty = threadIdx.y; + uint32_t pos = blockIdx.x; + uint32_t head_idx = ty; + + if (num_index_sets == 0) { + vec_t v; + v.fill(DTypeO(0.f)); + v.store(v_merged + (pos * num_heads + head_idx) * head_dim + tx * vec_size); + if (s_merged != nullptr) { + s_merged[pos * num_heads + head_idx] = -math::inf; + } + return; + } + + if (num_index_sets == 1) { + vec_t v; + v.cast_load(V + (pos * num_heads + head_idx) * head_dim + tx * vec_size); + v.store(v_merged + (pos * num_heads + head_idx) * head_dim + tx * vec_size); + if (s_merged != nullptr) { + s_merged[pos * num_heads + head_idx] = S[pos * num_heads + head_idx]; + } + return; + } + + state_t st; +#pragma unroll 2 + for (uint32_t iter = 0; iter < num_index_sets; ++iter) { + float s = S[(pos * num_index_sets + iter) * num_heads + head_idx]; + vec_t v; + v.cast_load(V + ((pos * num_index_sets + iter) * num_heads + head_idx) * head_dim + + tx * vec_size); + st.merge(v, s, 1); + } + + st.normalize(); + st.o.cast_store(v_merged + (pos * num_heads + head_idx) * head_dim + tx * vec_size); + if (s_merged != nullptr) { + s_merged[pos * num_heads + head_idx] = st.get_lse(); + } +} + +/*! + * \brief The CUDA kernel that merges self-attention states of a list of index sets, + * accelerated for larger number of index sets. + * \tparam vec_size The vector size used in the kernel. + * \tparam bdx The blockDim.x used in the kernel. + * \tparam bdy The blockDim.y used in the kernel. + * \tparam num_smem_stages The number of stages of shared memory used in the kernel. + * \tparam DTypeIn The data type of v. + * \tparam DTypeO The data type of v_merged. + * \param V The partial v of index sets. (n, num_index_sets, h, d) + * \param S The logsumexp value of index sets. (n, num_index_sets, h) + * \param v_merged The merged v of index sets union. (n, h, d) + * \param s_merged The merged logsumexp value of index sets union. (n, h) + * \param num_heads The number of heads of v. + * \param head_dim The dimension of each head. + * \note s are logsumexp values with base 2. + */ +template +__global__ void MergeStatesLargeNumIndexSetsKernel(DTypeIn* __restrict__ V, float* __restrict__ S, + DTypeO* __restrict__ v_merged, + float* __restrict__ s_merged, + uint32_t num_index_sets, uint32_t num_heads) { + uint32_t tx = threadIdx.x, ty = threadIdx.y; + uint32_t pos = blockIdx.x; + uint32_t head_idx = blockIdx.y; + state_t st; + constexpr uint32_t vec_bits = sizeof(DTypeIn) * vec_size * 8; + constexpr uint32_t head_dim = vec_size * bdx; + + extern __shared__ uint8_t smem[]; + DTypeIn* v_smem = (DTypeIn*)smem; + float* s_smem = (float*)(smem + num_smem_stages * bdy * head_dim * sizeof(DTypeIn)); + +#pragma unroll + for (uint32_t iter = 0; iter < num_smem_stages; ++iter) { + cp_async::pred_load( + v_smem + (iter * bdy + ty) * head_dim + tx * vec_size, + V + ((pos * num_index_sets + (iter * bdy + ty)) * num_heads + head_idx) * head_dim + + tx * vec_size, + (iter * bdy + ty) < num_index_sets); + cp_async::commit_group(); + } +#pragma unroll 4 + for (uint32_t iter = 0; iter < ceil_div(num_index_sets, bdy); ++iter) { + if (iter % bdx == 0) { + s_smem[ty * bdx + tx] = + iter * bdy + (ty * bdx + tx) < num_index_sets + ? S[(pos * num_index_sets + (iter * bdy + ty * bdx + tx)) * num_heads + head_idx] + : 0.f; + __syncthreads(); + } + cp_async::wait_group(); + __syncthreads(); + vec_t v; + v.cast_load(v_smem + ((iter % num_smem_stages) * bdy + ty) * head_dim + tx * vec_size); + if (iter * bdy + ty < num_index_sets) { + float s = s_smem[(iter % bdx) * bdy + ty]; + st.merge(v, s, 1); + } + __syncthreads(); + cp_async::pred_load( + v_smem + ((iter % num_smem_stages) * bdy + ty) * head_dim + tx * vec_size, + V + + ((pos * num_index_sets + ((iter + num_smem_stages) * bdy + ty)) * num_heads + + head_idx) * + head_dim + + tx * vec_size, + (iter + num_smem_stages) * bdy + ty < num_index_sets); + cp_async::commit_group(); + } + cp_async::wait_group<0>(); + __syncthreads(); + + st.normalize(); + threadblock_sync_state(st, v_smem, s_smem); + st.normalize(); + + st.o.cast_store(v_merged + (pos * num_heads + head_idx) * head_dim + tx * vec_size); + if (s_merged != nullptr) { + s_merged[pos * num_heads + head_idx] = st.get_lse(); + } +} + +/*! + * \brief The CUDA kernel to merge self-attention states of multiple index sets, the number of + * index sets at each position might vary. + * + * For CUDA graph support, the kernel can be built with a maximum sequence length and executed + * using a truncated, dynamic sequence length passed through `seq_len_ptr`. + * + * \tparam vec_size The vector size used in the kernel. + * \tparam bdx The blockDim.x used in the kernel. + * \tparam bdy The blockDim.y used in the kernel. + * \tparam num_smem_stages The number of stages of shared memory used in the kernel. + * \tparam DTypeIn The data type of v. + * \tparam DTypeO The data type of v_merged. + * \param V The partial v of index sets. (nnz, h, d) + * \param S The logsumexp value of index sets. (nnz, h) + * \param indptr The start offsets of each position in the variable length array. + * \param v_merged The merged v of index sets union. (n, h, d) + * \param s_merged The merged logsumexp value of index sets union. (n, h) + * \param max_seq_len The maximum sequence length supported by the kernel. + * \param seq_len_ptr The current sequence length (number of positions populated in indptr). + * \param num_heads The number of heads of v. + * \param head_dim The dimension of each head. + * \note s are logsumexp values with base 2. + */ +template +__global__ void PersistentVariableLengthMergeStatesKernel( + DTypeIn* __restrict__ V, float* __restrict__ S, IdType* indptr, DTypeO* __restrict__ v_merged, + float* __restrict__ s_merged, uint32_t max_seq_len, uint32_t* __restrict__ seq_len_ptr, + uint32_t num_heads) { + uint32_t tx = threadIdx.x, ty = threadIdx.y; + uint32_t cta_id = blockIdx.x; + uint32_t num_ctas = gridDim.x; + const uint32_t seq_len = seq_len_ptr ? *seq_len_ptr : max_seq_len; + uint32_t num_iters = ceil_div(seq_len * num_heads, num_ctas); + constexpr uint32_t vec_bits = sizeof(DTypeIn) * vec_size * 8; + constexpr uint32_t head_dim = vec_size * bdx; + extern __shared__ uint8_t smem[]; + DTypeIn* v_smem = (DTypeIn*)smem; + float* s_smem = (float*)(smem + num_smem_stages * bdy * head_dim * sizeof(DTypeIn)); + +#if (__CUDACC_VER_MAJOR__ >= 12 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.wait;"); +#endif + +#pragma unroll 1 + for (uint32_t i = cta_id; i < seq_len * num_heads; i += num_ctas) { + uint32_t pos = i / num_heads; + uint32_t head_idx = i % num_heads; + state_t st; + const uint32_t num_index_sets = indptr[pos + 1] - indptr[pos]; + + if (num_index_sets == 0) { + vec_t v; + v.fill(DTypeO(0.f)); + v.store(v_merged + (pos * num_heads + head_idx) * head_dim + tx * vec_size); + if (s_merged != nullptr) { + s_merged[pos * num_heads + head_idx] = -math::inf; + } + continue; + } + + if (num_index_sets == 1) { + vec_t v; + v.cast_load(V + (indptr[pos] * num_heads + head_idx) * head_dim + tx * vec_size); + v.store(v_merged + (pos * num_heads + head_idx) * head_dim + tx * vec_size); + if (s_merged != nullptr) { + s_merged[pos * num_heads + head_idx] = S[indptr[pos] * num_heads + head_idx]; + } + continue; + } + +#pragma unroll + for (uint32_t iter = 0; iter < num_smem_stages; ++iter) { + cp_async::pred_load( + v_smem + (iter * bdy + ty) * head_dim + tx * vec_size, + V + ((indptr[pos] + (iter * bdy + ty)) * num_heads + head_idx) * head_dim + tx * vec_size, + (iter * bdy + ty) < num_index_sets); + cp_async::commit_group(); + } +#pragma unroll 4 + for (uint32_t iter = 0; iter < ceil_div(num_index_sets, bdy); ++iter) { + if (iter % bdx == 0) { + s_smem[ty * bdx + tx] = + iter * bdy + (ty * bdx + tx) < num_index_sets + ? S[(indptr[pos] + (iter * bdy + ty * bdx + tx)) * num_heads + head_idx] + : 0.f; + __syncthreads(); + } + cp_async::wait_group(); + __syncthreads(); + vec_t v; + v.cast_load(v_smem + ((iter % num_smem_stages) * bdy + ty) * head_dim + tx * vec_size); + if (iter * bdy + ty < num_index_sets) { + float s = s_smem[(iter % bdx) * bdy + ty]; + st.merge(v, s, 1); + } + __syncthreads(); + cp_async::pred_load( + v_smem + ((iter % num_smem_stages) * bdy + ty) * head_dim + tx * vec_size, + V + + ((indptr[pos] + ((iter + num_smem_stages) * bdy + ty)) * num_heads + head_idx) * + head_dim + + tx * vec_size, + (iter + num_smem_stages) * bdy + ty < num_index_sets); + cp_async::commit_group(); + } + cp_async::wait_group<0>(); + __syncthreads(); + + st.normalize(); + threadblock_sync_state(st, v_smem, s_smem); + st.normalize(); + + st.o.cast_store(v_merged + (pos * num_heads + head_idx) * head_dim + tx * vec_size); + if (s_merged != nullptr) { + s_merged[pos * num_heads + head_idx] = st.get_lse(); + } + } +#if (__CUDACC_VER_MAJOR__ >= 12 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.launch_dependents;"); +#endif +} + +template +__global__ void PersistentVariableLengthAttentionSumKernel(DTypeIn* __restrict__ V, IdType* indptr, + DTypeO* __restrict__ v_sum, + uint32_t max_seq_len, + uint32_t* __restrict__ seq_len_ptr, + uint32_t num_heads) { + uint32_t tx = threadIdx.x, ty = threadIdx.y; + uint32_t cta_id = blockIdx.x; + uint32_t num_ctas = gridDim.x; + const uint32_t seq_len = seq_len_ptr ? *seq_len_ptr : max_seq_len; + uint32_t num_iters = ceil_div(seq_len * num_heads, num_ctas); + constexpr uint32_t vec_bits = sizeof(DTypeIn) * vec_size * 8; + constexpr uint32_t head_dim = vec_size * bdx; + extern __shared__ uint8_t smem[]; + DTypeIn* v_smem = (DTypeIn*)smem; + + vec_t v_sum_vec; +#if (__CUDACC_VER_MAJOR__ >= 12 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.wait;"); +#endif + +#pragma unroll 1 + for (uint32_t i = cta_id; i < seq_len * num_heads; i += num_ctas) { + uint32_t pos = i / num_heads; + uint32_t head_idx = i % num_heads; + const uint32_t num_index_sets = indptr[pos + 1] - indptr[pos]; + + if (num_index_sets == 0) { + vec_t v; + v.fill(DTypeO(0.f)); + v.store(v_sum + (pos * num_heads + head_idx) * head_dim + tx * vec_size); + continue; + } + + if (num_index_sets == 1) { + vec_t v; + v.cast_load(V + (indptr[pos] * num_heads + head_idx) * head_dim + tx * vec_size); + v.store(v_sum + (pos * num_heads + head_idx) * head_dim + tx * vec_size); + continue; + } + +#pragma unroll + for (uint32_t iter = 0; iter < num_smem_stages; ++iter) { + cp_async::pred_load( + v_smem + (iter * bdy + ty) * head_dim + tx * vec_size, + V + ((indptr[pos] + (iter * bdy + ty)) * num_heads + head_idx) * head_dim + tx * vec_size, + (iter * bdy + ty) < num_index_sets); + cp_async::commit_group(); + } +#pragma unroll 4 + for (uint32_t iter = 0; iter < ceil_div(num_index_sets, bdy); ++iter) { + cp_async::wait_group(); + __syncthreads(); + vec_t v; + v.cast_load(v_smem + ((iter % num_smem_stages) * bdy + ty) * head_dim + tx * vec_size); + if (iter * bdy + ty < num_index_sets) { +#pragma unroll + for (uint32_t i = 0; i < vec_size; ++i) { + v_sum_vec[i] += v[i]; + } + } + __syncthreads(); + cp_async::pred_load( + v_smem + ((iter % num_smem_stages) * bdy + ty) * head_dim + tx * vec_size, + V + + ((indptr[pos] + ((iter + num_smem_stages) * bdy + ty)) * num_heads + head_idx) * + head_dim + + tx * vec_size, + (iter + num_smem_stages) * bdy + ty < num_index_sets); + cp_async::commit_group(); + } + cp_async::wait_group<0>(); + __syncthreads(); + + threadblock_sum(v_sum_vec, v_smem); + + v_sum_vec.cast_store(v_sum + (pos * num_heads + head_idx) * head_dim + tx * vec_size); + } +#if (__CUDACC_VER_MAJOR__ >= 12 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.launch_dependents;"); +#endif +} + +/*! + * \brief Merge the self-attention state of two index sets A and B. + * \tparam DTypeIn The data type of v_a and v_b. + * \tparam DTypeO The data type of v_merged. + * \param v_a The partial v of index set A (n, h, d) + * \param s_a The logsumexp value of index set A. (n, h) + * \param v_b The partial v of index set B. (n, h, d) + * \param s_b The logsumexp value of index set B. (n, h) + * \param v_merged The merged v of index set A union B. (n, h, d) + * \param s_merged The merged logsumexp value of index set A union B. (n, h) + * \param seq_len The sequence length. + * \param num_heads The number of heads of v_a and v_b. + * \param head_dim The dimension of each head. + * \param stream The CUDA stream to execute the kernel. + * \return status Indicates whether CUDA calls are successful + * \note Both s_a and s_b are logsumexp values with base 2. + */ +template +cudaError_t MergeState(DTypeIn* v_a, float* s_a, DTypeIn* v_b, float* s_b, DTypeO* v_merged, + float* s_merged, uint32_t seq_len, uint32_t num_heads, uint32_t head_dim, + cudaStream_t stream = nullptr) { + DISPATCH_HEAD_DIM(head_dim, HEAD_DIM, { + constexpr uint32_t vec_size = std::max(16U / sizeof(DTypeIn), HEAD_DIM / 32U); + uint32_t bdx = HEAD_DIM / vec_size; + uint32_t bdy = num_heads; + dim3 nblks(seq_len); + dim3 nthrs(bdx, bdy); + auto kernel = MergeStateKernel; + void* args[] = {&v_a, &s_a, &v_b, &s_b, &v_merged, &s_merged, &num_heads, &head_dim}; + FLASHINFER_CUDA_CALL(cudaLaunchKernel((void*)kernel, nblks, nthrs, args, 0, stream)); + }); + return cudaSuccess; +} + +/*! + * \brief Merge the self-attention state with another state in place. + * \tparam DType The data type of v and v_other. + * \param v The partial v to be updated in-place. (n, h, d) + * \param s The logsumexp value to be updated in-place. (n, h) + * \param v_other The other v to be merged. (n, h, d) + * \param s_other The other logsumexp value to be merged. (n, h) + * \param seq_len The sequence length. + * \param num_heads The number of heads of v and v_other. + * \param head_dim The dimension of each head. + * \param mask Optional mask of whether to merge given sequences or not. (n) + * \param stream The CUDA stream to execute the kernel. + * \return status Indicates whether CUDA calls are successful + * \note Both s and s_other are logsumexp values with base 2. + */ +template +cudaError_t MergeStateInPlace(DType* v, float* s, DType* v_other, float* s_other, uint32_t seq_len, + uint32_t num_heads, uint32_t head_dim, uint8_t* mask = nullptr, + cudaStream_t stream = nullptr) { + DISPATCH_HEAD_DIM(head_dim, HEAD_DIM, { + constexpr uint32_t vec_size = std::max(16U / sizeof(DType), HEAD_DIM / 32U); + uint32_t bdx = HEAD_DIM / vec_size; + uint32_t bdy = num_heads; + dim3 nblks(seq_len); + dim3 nthrs(bdx, bdy); + auto kernel = MergeStateInPlaceKernel; + void* args[] = {&v, &s, &v_other, &s_other, &mask, &num_heads, &head_dim}; + FLASHINFER_CUDA_CALL(cudaLaunchKernel((void*)kernel, nblks, nthrs, args, 0, stream)); + }); + return cudaSuccess; +} + +/*! + * \brief Merge self-attention states of a list of index sets. + * \tparam DTypeIn The data type of v. + * \tparam DTypeO The data type of v_merged. + * \param v The partial v of index sets. (n, num_index_sets, h, d) + * \param s The logsumexp value of index sets. (n, num_index_sets, h) + * \param v_merged The merged v of index sets union. (n, h, d) + * \param s_merged The merged logsumexp value of index sets union. (n, h) + * \param num_index_sets The number of index sets. + * \param seq_len The sequence length. + * \param num_heads The number of heads of v. + * \param head_dim The dimension of each head. + * \param stream The CUDA stream to execute the kernel. + * \return status Indicates whether CUDA calls are successful + * \note s are logsumexp values with base 2. + */ +template +cudaError_t MergeStates(DTypeIn* v, float* s, DTypeO* v_merged, float* s_merged, + uint32_t num_index_sets, uint32_t seq_len, uint32_t num_heads, + uint32_t head_dim, cudaStream_t stream = nullptr) { + DISPATCH_HEAD_DIM(head_dim, HEAD_DIM, { + constexpr uint32_t vec_size = std::max(16U / sizeof(DTypeIn), HEAD_DIM / 32U); + constexpr uint32_t bdx = HEAD_DIM / vec_size; + if (num_index_sets >= seq_len) { + constexpr uint32_t num_threads = 128; + constexpr uint32_t bdy = num_threads / bdx; + dim3 nblks(seq_len, num_heads); + dim3 nthrs(bdx, bdy); + constexpr uint32_t num_smem_stages = 4; + auto kernel = + MergeStatesLargeNumIndexSetsKernel; + void* args[] = {&v, &s, &v_merged, &s_merged, &num_index_sets, &num_heads}; + uint32_t smem_size = + num_smem_stages * bdy * head_dim * sizeof(DTypeIn) + num_threads * sizeof(float); + FLASHINFER_CUDA_CALL( + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + FLASHINFER_CUDA_CALL(cudaLaunchKernel((void*)kernel, nblks, nthrs, args, smem_size, stream)); + } else { + uint32_t bdy = num_heads; + dim3 nblks(seq_len); + dim3 nthrs(bdx, bdy); + auto kernel = MergeStatesKernel; + void* args[] = {&v, &s, &v_merged, &s_merged, &num_index_sets, &num_heads, &head_dim}; + FLASHINFER_CUDA_CALL(cudaLaunchKernel((void*)kernel, nblks, nthrs, args, 0, stream)); + } + }); + return cudaSuccess; +} + +template +cudaError_t AttentionSum(DTypeIn* v, DTypeO* v_sum, uint32_t num_index_sets, uint32_t seq_len, + uint32_t num_heads, uint32_t head_dim, cudaStream_t stream = nullptr) { + DISPATCH_HEAD_DIM(head_dim, HEAD_DIM, { + constexpr uint32_t vec_size = std::max(16U / sizeof(DTypeIn), HEAD_DIM / 32U); + constexpr uint32_t bdx = HEAD_DIM / vec_size; + uint32_t bdy = num_heads; + dim3 nblks(seq_len); + dim3 nthrs(bdx, bdy); + auto kernel = AttentionSumKernel; + void* args[] = {&v, &v_sum, &num_index_sets, &num_heads, &head_dim}; + FLASHINFER_CUDA_CALL(cudaLaunchKernel((void*)kernel, nblks, nthrs, args, 0, stream)); + }); + return cudaSuccess; +} + +template +cudaError_t VariableLengthMergeStates(DTypeIn* v, float* s, IdType* indptr, DTypeO* v_merged, + float* s_merged, uint32_t max_seq_len, uint32_t* seq_len, + uint32_t num_heads, uint32_t head_dim, bool enable_pdl, + cudaStream_t stream = nullptr) { + int dev_id = 0; + int num_sms = 0; + int num_blocks_per_sm = 0; + FLASHINFER_CUDA_CALL(cudaGetDevice(&dev_id)); + FLASHINFER_CUDA_CALL(cudaDeviceGetAttribute(&num_sms, cudaDevAttrMultiProcessorCount, dev_id)); + + DISPATCH_HEAD_DIM(head_dim, HEAD_DIM, { + constexpr uint32_t vec_size = std::max(16U / sizeof(DTypeIn), HEAD_DIM / 32U); + constexpr uint32_t bdx = HEAD_DIM / vec_size; + constexpr uint32_t num_threads = 128; + constexpr uint32_t bdy = num_threads / bdx; + constexpr uint32_t num_smem_stages = 4; + uint32_t smem_size = + num_smem_stages * bdy * head_dim * sizeof(DTypeIn) + num_threads * sizeof(float); + auto kernel = PersistentVariableLengthMergeStatesKernel; + FLASHINFER_CUDA_CALL(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&num_blocks_per_sm, kernel, + num_threads, smem_size)); + num_blocks_per_sm = min(num_blocks_per_sm, ceil_div(max_seq_len * num_heads, num_sms)); + + dim3 nblks(num_sms * num_blocks_per_sm); + dim3 nthrs(bdx, bdy); + void* args[] = {&v, &s, &indptr, &v_merged, &s_merged, &max_seq_len, &seq_len, &num_heads}; + FLASHINFER_CUDA_CALL( + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + + // PDL launch + if (enable_pdl) { + cudaLaunchAttribute attribute[1]; + attribute[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attribute[0].val.programmaticStreamSerializationAllowed = 1; + cudaLaunchConfig_t config; + config.attrs = attribute; + config.numAttrs = 1; + config.gridDim = nblks; + config.blockDim = nthrs; + config.dynamicSmemBytes = smem_size; + config.stream = stream; + FLASHINFER_CUDA_CALL(cudaLaunchKernelEx(&config, kernel, v, s, indptr, v_merged, s_merged, + max_seq_len, seq_len, num_heads)); + } else { + FLASHINFER_CUDA_CALL(cudaLaunchKernel((void*)kernel, nblks, nthrs, args, smem_size, stream)); + } + }); + return cudaSuccess; +} + +template +cudaError_t VariableLengthAttentionSum(DTypeIn* v, IdType* indptr, DTypeO* v_sum, + uint32_t max_seq_len, uint32_t* seq_len, uint32_t num_heads, + uint32_t head_dim, bool enable_pdl, + cudaStream_t stream = nullptr) { + int dev_id = 0; + int num_sms = 0; + int num_blocks_per_sm = 0; + FLASHINFER_CUDA_CALL(cudaGetDevice(&dev_id)); + FLASHINFER_CUDA_CALL(cudaDeviceGetAttribute(&num_sms, cudaDevAttrMultiProcessorCount, dev_id)); + + DISPATCH_HEAD_DIM(head_dim, HEAD_DIM, { + constexpr uint32_t vec_size = std::max(16U / sizeof(DTypeIn), HEAD_DIM / 32U); + constexpr uint32_t bdx = HEAD_DIM / vec_size; + constexpr uint32_t num_threads = 128; + constexpr uint32_t bdy = num_threads / bdx; + constexpr uint32_t num_smem_stages = 4; + uint32_t smem_size = num_smem_stages * bdy * head_dim * sizeof(DTypeIn); + auto kernel = PersistentVariableLengthAttentionSumKernel; + FLASHINFER_CUDA_CALL(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&num_blocks_per_sm, kernel, + num_threads, smem_size)); + num_blocks_per_sm = min(num_blocks_per_sm, ceil_div(max_seq_len * num_heads, num_sms)); + + dim3 nblks(num_sms * num_blocks_per_sm); + dim3 nthrs(bdx, bdy); + void* args[] = {&v, &indptr, &v_sum, &max_seq_len, &seq_len, &num_heads}; + FLASHINFER_CUDA_CALL( + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + + if (enable_pdl) { + // PDL launch + cudaLaunchAttribute attribute[1]; + attribute[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attribute[0].val.programmaticStreamSerializationAllowed = 1; + cudaLaunchConfig_t config; + config.attrs = attribute; + config.numAttrs = 1; + config.gridDim = nblks; + config.blockDim = nthrs; + config.dynamicSmemBytes = smem_size; + config.stream = stream; + FLASHINFER_CUDA_CALL( + cudaLaunchKernelEx(&config, kernel, v, indptr, v_sum, max_seq_len, seq_len, num_heads)); + } else { + FLASHINFER_CUDA_CALL(cudaLaunchKernel((void*)kernel, nblks, nthrs, args, smem_size, stream)); + } + }); + return cudaSuccess; +} + +} // namespace flashinfer + +#endif // FLASHINFER_CASCADE_CUH_ diff --git a/include/flashinfer/attention/cutlass_mla.cuh b/include/flashinfer/attention/cutlass_mla.cuh new file mode 100644 index 0000000000000000000000000000000000000000..cd0b42413d06fa4c016bed3b15d9cd469a88dc99 --- /dev/null +++ b/include/flashinfer/attention/cutlass_mla.cuh @@ -0,0 +1,153 @@ +/* + * Copyright (c) 2024 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef FLASHINFER_ATTENTION_CUTLASS_MLA_CUH_ +#define FLASHINFER_ATTENTION_CUTLASS_MLA_CUH_ +#include + +#include "../cutlass_utils.cuh" +#include "../exception.h" +#include "cutlass/kernel_hardware_info.h" + +// From 3rdparty/cutlass/examples/77_blackwell_fmha +#include "blackwell/device/sm100_mla.hpp" +#include "blackwell/kernel/sm100_mla_tile_scheduler.hpp" + +namespace flashinfer { + +namespace attention { + +using namespace cute; +using namespace cutlass::fmha::kernel; + +template +struct IsPersistent { + static const bool value = v; +}; + +template > +struct MlaSm100 { + using Element = T; + using ElementAcc = float; + using ElementOut = T; + + using TileShape = Shape<_128, _128, Shape<_512, _64>>; + using TileShapeH = cute::tuple_element_t<0, TileShape>; + using TileShapeD = cute::tuple_element_t<2, TileShape>; + + // H K (D_latent D_rope) B + using ProblemShape = cute::tuple; + + using StrideQ = cute::tuple; // H D B + using StrideK = cute::tuple; // K D B + using StrideO = StrideK; // H D B + using StrideLSE = cute::tuple<_1, int>; // H B + + using TileScheduler = + std::conditional_t; + + using FmhaKernel = cutlass::fmha::kernel::Sm100FmhaMlaKernelTmaWarpspecialized< + TileShape, Element, ElementAcc, ElementOut, ElementAcc, TileScheduler, /*kIsCpAsync=*/true>; + using Fmha = cutlass::fmha::device::MLA; +}; + +template +typename T::Fmha::Arguments args_from_options(void* out_ptr, void* lse_ptr, void* q_absorbed_ptr, + void* ckv_kpe_cache_ptr, void* seq_lens_ptr, + void* page_table_ptr, int batches, + int page_count_per_seq, int page_count_total, + int page_size, int device_index) { + cutlass::KernelHardwareInfo hw_info; + hw_info.device_id = device_index; + hw_info.sm_count = + cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id); + + int max_seq_len = page_size * page_count_per_seq; + using TileShapeH = typename T::TileShapeH; + using TileShapeD = typename T::TileShapeD; + auto problem_shape = cute::make_tuple(TileShapeH{}, max_seq_len, TileShapeD{}, batches); + + auto [H, K, D, B] = problem_shape; + auto [D_latent, D_rope] = D; + + // the scale is based on the non-absorbed sizes, change as appropriate + // we can't determine this parameter from the info we have, it's an input + int D_non_latent = 128; + float scale = 1.0 / sqrt(1.0 * (D_non_latent + D_rope)); + + using StrideQ = typename T::StrideQ; + using StrideK = typename T::StrideK; + using StrideO = typename T::StrideO; + using StrideLSE = typename T::StrideLSE; + + StrideQ stride_Q = cute::make_tuple(static_cast(0 + D_latent + D_rope), _1{}, + static_cast(H * (0 + D_latent + D_rope))); + StrideK stride_C = cute::make_tuple(static_cast(0 + D_latent + D_rope), _1{}, + static_cast(page_size * (D_latent + D_rope))); + StrideLSE stride_PT = cute::make_stride(_1{}, page_count_per_seq); + StrideLSE stride_LSE = cute::make_tuple(_1{}, 0 + H); + StrideO stride_O = cute::make_tuple(static_cast(0 + D_latent), _1{}, + static_cast(0 + H * D_latent)); + + using Element = typename T::Element; + using ElementOut = typename T::ElementOut; + using ElementAcc = typename T::ElementAcc; + auto Q_ptr = reinterpret_cast(q_absorbed_ptr); + auto C_ptr = reinterpret_cast(ckv_kpe_cache_ptr); + typename T::Fmha::Arguments arguments{ + problem_shape, + {scale, Q_ptr, stride_Q, Q_ptr + D_latent, stride_Q, C_ptr, stride_C, C_ptr + D_latent, + stride_C, reinterpret_cast(seq_lens_ptr), reinterpret_cast(page_table_ptr), + stride_PT, page_count_total, page_size}, + {reinterpret_cast(out_ptr), stride_O, + // static_cast(lse.data_ptr()), stride_LSE}, + static_cast(nullptr), stride_LSE}, + hw_info, + -1, // split_kv + nullptr, // is_var_split_kv=false + }; + // TODO(kaixih@nvidia): When split_kv=-1 and is_var_split_kv=false, we compute + // split_kv automatically based on batch size and sequence length to balance + // workload across available SMs. Consider using var_split_kv for manual + // control if needed. + T::Fmha::set_split_kv(arguments); + return arguments; +} + +template +cudaError_t runMla(void* workspace_ptr, void* out_ptr, void* lse_ptr, void* q_absorbed_ptr, + void* ckv_kpe_cache_ptr, void* seq_lens_ptr, void* page_table_ptr, int batches, + int page_count_per_seq, int page_count_total, int page_size, int device_index, + cudaStream_t stream) { + using MlaSm100Type = MlaSm100; + typename MlaSm100Type::Fmha fmha; + auto arguments = args_from_options( + out_ptr, lse_ptr, q_absorbed_ptr, ckv_kpe_cache_ptr, seq_lens_ptr, page_table_ptr, batches, + page_count_per_seq, page_count_total, page_size, device_index); + + CUTLASS_CHECK(fmha.can_implement(arguments)); + + CUTLASS_CHECK(fmha.initialize(arguments, workspace_ptr, stream)); + + CUTLASS_CHECK(fmha.run(arguments, workspace_ptr, stream)); + + return cudaSuccess; +} + +} // namespace attention + +} // namespace flashinfer +#endif // FLASHINFER_ATTENTION_CUTLASS_MLA_CUH_ diff --git a/include/flashinfer/attention/decode.cuh b/include/flashinfer/attention/decode.cuh new file mode 100644 index 0000000000000000000000000000000000000000..ed173fb618d10f19a80e9219be31c219b016e669 --- /dev/null +++ b/include/flashinfer/attention/decode.cuh @@ -0,0 +1,1173 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef FLASHINFER_DECODE_CUH_ +#define FLASHINFER_DECODE_CUH_ +#include +#include +#include +#include +#include + +#include + +#include "../cp_async.cuh" +#include "../math.cuh" +#include "../pos_enc.cuh" +#include "../utils.cuh" +#include "../vec_dtypes.cuh" +#include "cascade.cuh" +#include "state.cuh" + +namespace flashinfer { + +DEFINE_HAS_MEMBER(decode_maybe_q_rope_offset) + +namespace cg = cooperative_groups; +using cp_async::PrefetchMode; +using cp_async::SharedMemFillMode; + +namespace { + +/*! + * \brief Load k tile from smem and compute qk + * \tparam pos_encoding_mode The positional encoding mode used in the kernel + * \tparam head_dim A template integer indicates the head dimension + * \tparam vec_size A template integer indicates the vector size + * \tparam bdx A template integer indicates the block size in x dimension + * \tparam tile_size A template integer indicates the tile size per (bdx * bdy) threads. + * \tparam T A template type indicates the input data type + * \param smem A pointer to the start of shared memory + * \param q_vec A vector of float indicates the thread-local query vector + * \param freq A vector of float indicates the thread-local rope frequency + * \param kv_shared_offset An array of uint32_t indicates the k/v tiles offset + * in shared memory of different pipeline stages + * \param kv_idx A integer indicates the thread-local kv position in kv-cache + * \param compute_stage_idx A integer indicates the compute stage index in the pipeline + * \param s A float indicates the thread-local result of qk + * \param st The self-attention state to be updated + */ +template +__device__ __forceinline__ void compute_qk( + const Params& params, AttentionVariant variant, const uint32_t batch_idx, const T* smem, + const vec_t& q_vec, const vec_t& freq, uint32_t kv_idx_base, + uint32_t iter_base, uint32_t iter_bound, uint32_t qo_head_idx, uint32_t kv_head_idx, float* s, + state_t& st, const uint32_t tx, const uint32_t ty, const uint32_t tz) { + float m_prev = st.m; +#pragma unroll + for (uint32_t j = 0; j < tile_size; ++j) { + vec_t k_vec; + if constexpr (pos_encoding_mode == PosEncodingMode::kRoPELlama) { + // apply rotary embedding for all rows in k matrix of kv-cache + k_vec = vec_apply_llama_rope(smem + j * bdx * vec_size, freq, + kv_idx_base + tz * tile_size + j); + } else { + // do not apply rotary embedding + k_vec.cast_load(smem + (j * bdx + tx) * vec_size); + } + s[j] = 0.f; +#pragma unroll + for (uint32_t i = 0; i < vec_size; ++i) { + s[j] += q_vec[i] * k_vec[i]; + } +#pragma unroll + for (uint32_t offset = bdx / 2; offset > 0; offset /= 2) { + s[j] += math::shfl_xor_sync(s[j], offset); + } + const uint32_t pos = kv_idx_base + tz * tile_size + j; + s[j] = variant.LogitsTransform(params, s[j], batch_idx, /*qo_idx=*/0, /*kv_idx=*/pos, + qo_head_idx, kv_head_idx); + if constexpr (variant.use_softmax) { + s[j] *= variant.sm_scale_log2; + } + + bool mask = variant.LogitsMask(params, batch_idx, /*qo_idx=*/0, /*kv_idx=*/pos, qo_head_idx, + kv_head_idx); + s[j] = (iter_base + tz * tile_size + j < iter_bound && mask) ? s[j] : -math::inf; + st.m = max(st.m, s[j]); + } + + if constexpr (variant.use_softmax) { + float o_scale = math::ptx_exp2(m_prev - st.m); + st.d *= o_scale; +#pragma unroll + for (uint32_t j = 0; j < tile_size; ++j) { + s[j] = math::ptx_exp2(s[j] - st.m); + st.d += s[j]; + } +#pragma unroll + for (uint32_t i = 0; i < vec_size; ++i) { + st.o[i] = st.o[i] * o_scale; + } + } +} + +/*! + * \brief Load v tile from shared memory and update local state + * \tparam vec_size A template integer indicates the vector size + * \tparam bdx A template integer indicates the block size in x dimension + * \tparam tile_size A template integer indicates the tile size per (bdx * bdy) threads. + * \tparam T A template type indicates the input data type + * \param smem A pointer to the start of shared memory + * \param s A float indicates the pre-softmax attention score + * \param kv_shared_offset An array of uint32_t indicates the k/v tiles offset + * in shared memory of different pipeline stages + * \param compute_stage_idx A integer indicates the compute stage index in the pipeline + * \param st The flashattention state to be updated + */ +template +__device__ __forceinline__ void update_local_state(const T* smem, const float* s, + uint32_t compute_stage_idx, + state_t& st, uint32_t tx) { +#pragma unroll + for (uint32_t j = 0; j < tile_size; ++j) { + vec_t v_vec; + v_vec.cast_load(smem + (j * bdx + tx) * vec_size); +#pragma unroll + for (uint32_t i = 0; i < vec_size; ++i) { + st.o[i] = st.o[i] + s[j] * v_vec[i]; + } + } +} + +/*! + * \brief Synchronize the state of all warps inside a threadblock. + * \tparam vec_size A template integer indicates the vector size + * \tparam bdx A template integer indicates the block size in x dimension + * \tparam bdy A template integer indicates the block size in y dimension + * \param st The warp local state + * \param smem The pointer to shared memory buffer for o + * \param smem_md The pointer to shared memory buffer for m/d + */ +template +__device__ __forceinline__ void sync_state(AttentionVariant variant, state_t& st, + float* smem, float* smem_md, const uint32_t tx, + const uint32_t ty, const uint32_t tz) { + if constexpr (bdz > 1) { + constexpr uint32_t head_dim = bdx * vec_size; + auto block = cg::this_thread_block(); + st.o.store(smem + (tz * bdy + ty) * head_dim + tx * vec_size); + if constexpr (variant.use_softmax) { + smem_md[(tz * bdy + ty) * 2] = st.m; + smem_md[(tz * bdy + ty) * 2 + 1] = st.d; + block.sync(); + st.init(); +#pragma unroll + for (uint32_t j = 0; j < bdz; ++j) { + float mz = smem_md[(j * bdy + ty) * 2], dz = smem_md[(j * bdy + ty) * 2 + 1]; + vec_t oz; + oz.load(smem + (j * bdy + ty) * head_dim + tx * vec_size); + st.merge(oz, mz, dz); + } + } else { + block.sync(); + st.init(); +#pragma unroll + for (uint32_t j = 0; j < bdz; ++j) { + vec_t oz; + oz.load(smem + (j * bdy + ty) * head_dim + tx * vec_size); +#pragma unroll + for (uint32_t i = 0; i < vec_size; ++i) { + st.o[i] += oz[i]; + } + } + } + } +} + +} // namespace + +/*! + * \brief FlashAttention decoding cuda kernel with kv-cache for a single request + * \tparam pos_encoding_mode The positional encoding mode + * \tparam vec_size A template integer indicates the vector size + * \tparam bdx A template integer indicates the block size in x dimension + * \tparam bdy A template integer indicates the block size in y dimension + * \tparam DTypeQ A template type indicates the query data type + * \tparam DTypeKV A template type indicates the key-value data type + * \tparam DTypeO A template type indicates the output data type + * \param q [num_qo_heads, head_dim] The query matrix + * \param k [seq_len, num_kv_heads, head_dim] The key matrix in kv-cache + * \param v [seq_len, num_kv_heads, head_dim] The value matrix in kv-cache + * \param o [num_qo_heads, head_dim] The output matrix + * \param head_dim A integer indicates the head dimension + * \param rope_rcp_scale A floating number indicate the reciprocal + * of scaling ratio used in PI(Position Interpolation) for RoPE (Rotary + * Positional Embeddings) + * \param rope_rcp_theta A floating number indicate the reciprocal + * of "theta" used in RoPE (Rotary Positional Embeddings) + * \param kv_chunk_size A integer indicates the kv-chunk size + */ +template +__global__ void SingleDecodeWithKVCacheKernel(const __grid_constant__ Params params) { + using DTypeQ = typename Params::DTypeQ; + using DTypeKV = typename Params::DTypeKV; + using DTypeO = typename Params::DTypeO; + const DTypeQ* q = params.q; + const DTypeKV* k = params.k; + const DTypeKV* v = params.v; + const uint32_t q_stride_n = params.q_stride_n; + const uint32_t q_stride_h = params.q_stride_h; + const uint32_t kv_stride_n = params.kv_stride_n; + const uint32_t kv_stride_h = params.kv_stride_h; + DTypeO* o = params.o; + float* lse = params.lse; + uint32_t kv_chunk_size = params.kv_chunk_size; + + auto block = cg::this_thread_block(); + auto grid = cg::this_grid(); + + constexpr uint32_t head_dim = bdx * vec_size; + uint32_t kv_head_idx = blockIdx.y; + uint32_t qo_head_idx = kv_head_idx * bdy + threadIdx.y; + uint32_t kv_chunk_idx = blockIdx.x; + uint32_t num_qo_heads = params.num_qo_heads; + + extern __shared__ uint8_t smem[]; + AttentionVariant variant(params, /*batch_idx=*/0, smem); + const uint32_t seq_len = variant.kv_len; + DTypeKV* k_smem = (DTypeKV*)smem; + DTypeKV* v_smem = (DTypeKV*)(smem + num_stages_smem * bdy * tile_size_per_bdx * bdz * head_dim * + sizeof(DTypeKV)); + float* smem_md = (float*)(smem + 2 * num_stages_smem * bdy * tile_size_per_bdx * bdz * head_dim * + sizeof(DTypeKV)); + + uint32_t tx = threadIdx.x, ty = threadIdx.y, tz = threadIdx.z; + vec_t q_vec; + vec_t freq; + if constexpr (pos_encoding_mode == PosEncodingMode::kRoPELlama) { + const float rope_rcp_scale = params.rope_rcp_scale; + const float rope_rcp_theta = params.rope_rcp_theta; + +#pragma unroll + for (uint32_t i = 0; i < vec_size; ++i) { + freq[i] = rope_rcp_scale * + __powf(rope_rcp_theta, + float(2 * ((tx * vec_size + i) % (head_dim / 2))) / float(head_dim)); + } + + // apply rotary embedding to q matrix + q_vec = vec_apply_llama_rope(q + qo_head_idx * q_stride_h, freq, seq_len - 1); + } else { + // do not apply rotary embedding to q matrix + q_vec.cast_load(q + qo_head_idx * q_stride_h + tx * vec_size); + } + block.sync(); + + uint32_t chunk_start = kv_chunk_idx * kv_chunk_size; + kv_chunk_size = min(kv_chunk_size, seq_len - chunk_start); + uint32_t chunk_end = chunk_start + kv_chunk_size; + + // preload k tiles and v tiles + uint32_t producer_kv_idx_base = chunk_start; + constexpr uint32_t vec_bits = sizeof(DTypeKV) * vec_size * 8; +#pragma unroll + for (uint32_t iter = 0; iter < num_stages_smem; ++iter) { + for (uint32_t j = 0; j < tile_size_per_bdx; ++j) { + cp_async::pred_load( + k_smem + (((iter * bdz + tz) * bdy + ty) * tile_size_per_bdx + j) * head_dim + + tx * vec_size, + k + (producer_kv_idx_base + (tz * bdy + ty) * tile_size_per_bdx + j) * kv_stride_n + + kv_head_idx * kv_stride_h + tx * vec_size, + producer_kv_idx_base + (tz * bdy + ty) * tile_size_per_bdx + j < chunk_end); + } + cp_async::commit_group(); + for (uint32_t j = 0; j < tile_size_per_bdx; ++j) { + cp_async::pred_load( + v_smem + (((iter * bdz + tz) * bdy + ty) * tile_size_per_bdx + j) * head_dim + + tx * vec_size, + v + (producer_kv_idx_base + (tz * bdy + ty) * tile_size_per_bdx + j) * kv_stride_n + + kv_head_idx * kv_stride_h + tx * vec_size, + producer_kv_idx_base + (tz * bdy + ty) * tile_size_per_bdx + j < chunk_end); + } + cp_async::commit_group(); + producer_kv_idx_base += bdy * bdz * tile_size_per_bdx; + } + + // pipelining k/v tiles loading and state updating + uint32_t consumer_kv_idx_base = chunk_start, stage_idx = 0; + state_t st_local; + float s[bdy * tile_size_per_bdx]; + +#pragma unroll 2 + for (uint32_t iter = 0; iter < ceil_div(kv_chunk_size, tile_size_per_bdx * bdy * bdz); ++iter) { + // compute qk + cp_async::wait_group<2 * num_stages_smem - 1>(); + block.sync(); + compute_qk( + params, variant, /*batch_idx=*/0, + k_smem + (stage_idx * bdz + tz) * bdy * tile_size_per_bdx * head_dim, q_vec, freq, + consumer_kv_idx_base, iter * bdy * tile_size_per_bdx * bdz, kv_chunk_size, qo_head_idx, + kv_head_idx, s, st_local, tx, ty, tz); + block.sync(); + // load k + for (uint32_t j = 0; j < tile_size_per_bdx; ++j) { + cp_async::pred_load( + k_smem + (((stage_idx * bdz + tz) * bdy + ty) * tile_size_per_bdx + j) * head_dim + + tx * vec_size, + k + (producer_kv_idx_base + (tz * bdy + ty) * tile_size_per_bdx + j) * kv_stride_n + + kv_head_idx * kv_stride_h + tx * vec_size, + producer_kv_idx_base + (tz * bdy + ty) * tile_size_per_bdx + j < chunk_end); + } + cp_async::commit_group(); + + // update m/d/o state + cp_async::wait_group<2 * num_stages_smem - 1>(); + block.sync(); + update_local_state( + v_smem + (stage_idx * bdz + tz) * bdy * tile_size_per_bdx * head_dim, s, stage_idx, + st_local, tx); + block.sync(); + + // load v + for (uint32_t j = 0; j < tile_size_per_bdx; ++j) { + cp_async::pred_load( + v_smem + (((stage_idx * bdz + tz) * bdy + ty) * tile_size_per_bdx + j) * head_dim + + tx * vec_size, + v + (producer_kv_idx_base + (tz * bdy + ty) * tile_size_per_bdx + j) * kv_stride_n + + kv_head_idx * kv_stride_h + tx * vec_size, + producer_kv_idx_base + (tz * bdy + ty) * tile_size_per_bdx + j < chunk_end); + } + cp_async::commit_group(); + + stage_idx = (stage_idx + 1) % num_stages_smem; + producer_kv_idx_base += tile_size_per_bdx * bdy * bdz; + consumer_kv_idx_base += tile_size_per_bdx * bdy * bdz; + } + cp_async::wait_group<0>(); + block.sync(); + + // sync local state of all warps inside a threadblock + sync_state(variant, st_local, reinterpret_cast(smem), smem_md, + tx, ty, tz); + if constexpr (variant.use_softmax) { + st_local.normalize(); + } + + st_local.o.cast_store(o + (kv_chunk_idx * num_qo_heads + qo_head_idx) * head_dim + tx * vec_size); + if (lse != nullptr) { + lse[kv_chunk_idx * num_qo_heads + qo_head_idx] = st_local.get_lse(); + } +} + +/*! + * \brief FlashAttention decoding cuda kernel with paged kv-cache for multiple requests + * \tparam pos_encoding_mode The positional encoding mode + * \tparam vec_size A template integer indicates the vector size + * \tparam bdx A template integer indicates the block size in x dimension + * \tparam bdy A template integer indicates the block size in y dimension + * \tparam bdz A template integer indicates the block size in z dimension + * \tparam DTypeQ A template type indicates the query data type + * \tparam DTypeKV A template type indicates the key-value data type + * \tparam DTypeO A template type indicates the output data type + * \tparam IdType A template type indicates the index data type + * \param q [batch_size, num_qo_heads, head_dim] The query matrix + * \param paged_kv The paged kv-cache data structure + * \param o [num_qo_heads, head_dim] The output matrix + * \param tmp Used-allocated temporary buffer + * \param lse The logsumexp values + * \param sm_scale A float indicates the scale applied to pre-softmax logits + * \param rope_rcp_scale A floating number indicate the reciprocal + * of scaling ratio used in PI(Position Interpolation) for RoPE (Rotary + * Positional Embeddings) + * \param rope_rcp_theta A floating number indicate the reciprocal + * of "theta" used in RoPE (Rotary Positional Embeddings) + */ +template +__device__ __inline__ void BatchDecodeWithPagedKVCacheDevice(const Params& params, uint8_t smem[], + const uint32_t bx = blockIdx.x, + const uint32_t by = blockIdx.y, + const uint32_t tx = threadIdx.x, + const uint32_t ty = threadIdx.y, + const uint32_t tz = threadIdx.z) { + auto block = cg::this_thread_block(); + using DTypeQ = typename Params::DTypeQ; + using DTypeKV = typename Params::DTypeKV; + using DTypeO = typename Params::DTypeO; + using IdType = typename Params::IdType; + const DTypeQ* q = params.q; + DTypeO* o = params.o; + float* lse = params.lse; + const auto paged_kv = params.paged_kv; + const bool* block_valid_mask = params.block_valid_mask; + const uint32_t padded_batch_size = params.padded_batch_size; + const uint32_t num_qo_heads = params.num_qo_heads; + const bool partition_kv = params.partition_kv; + + constexpr uint32_t head_dim = bdx * vec_size; + const uint32_t batch_idx = params.request_indices[bx]; + const uint32_t kv_tile_idx = params.kv_tile_indices[bx]; + const uint32_t kv_head_idx = by; + const uint32_t qo_head_idx = kv_head_idx * bdy + ty; + // NOTE(Zihao): when CUDAGraph is enabled, we will launch more blocks than + // the actual batch size, so we need to check if the current batch is valid + if (block_valid_mask && !block_valid_mask[bx]) return; + const uint32_t kv_chunk_size = *(params.kv_chunk_size_ptr); + const uint32_t kv_len = paged_kv.get_length(batch_idx); + const uint32_t max_chunk_size = partition_kv ? kv_chunk_size : kv_len; + const uint32_t chunk_start = partition_kv ? kv_tile_idx * max_chunk_size : 0; + const uint32_t chunk_end = + partition_kv ? min((kv_tile_idx + 1) * max_chunk_size, kv_len) : kv_len; + const uint32_t chunk_size = chunk_end - chunk_start; + + AttentionVariant variant(params, batch_idx, smem); + DTypeKV* k_smem = (DTypeKV*)smem; + DTypeKV* v_smem = (DTypeKV*)(smem + num_stages_smem * tile_size_per_bdx * bdy * bdz * head_dim * + sizeof(DTypeKV)); + size_t* kv_offset_smem = (size_t*)(smem + 2 * num_stages_smem * tile_size_per_bdx * bdy * bdz * + head_dim * sizeof(DTypeKV)); + float* smem_md = (float*)(smem + 2 * num_stages_smem * tile_size_per_bdx * bdy * bdz * head_dim * + sizeof(DTypeKV)); + + vec_t q_vec; + vec_t freq; + const uint32_t q_stride_n = params.q_stride_n; + const uint32_t q_stride_h = params.q_stride_h; + if constexpr (POS_ENCODING_MODE == PosEncodingMode::kRoPELlama) { + const IdType* q_rope_offset = nullptr; + if constexpr (has_decode_maybe_q_rope_offset_v) { + q_rope_offset = params.decode_maybe_q_rope_offset; + } + int32_t q_rope_offset_val = q_rope_offset == nullptr ? (kv_len - 1) : q_rope_offset[batch_idx]; + const float rope_rcp_scale = params.rope_rcp_scale; + const float rope_rcp_theta = params.rope_rcp_theta; + +#pragma unroll + for (uint32_t i = 0; i < vec_size; ++i) { + freq[i] = rope_rcp_scale * + __powf(rope_rcp_theta, + float(2 * ((tx * vec_size + i) % (head_dim / 2))) / float(head_dim)); + } +#if (__CUDACC_VER_MAJOR__ >= 12 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.wait;"); +#endif + // apply rotary embedding to q matrix + q_vec = vec_apply_llama_rope( + q + batch_idx * q_stride_n + qo_head_idx * q_stride_h, freq, q_rope_offset_val); + } else { +// do not apply rotary embedding to q matrix +#if (__CUDACC_VER_MAJOR__ >= 12 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.wait;"); +#endif + q_vec.cast_load(q + batch_idx * q_stride_n + qo_head_idx * q_stride_h + tx * vec_size); + } + + // preload k/v tiles + uint32_t stage_idx = 0; + constexpr uint32_t vec_bits = sizeof(DTypeKV) * vec_size * 8; + const IdType last_indptr = paged_kv.indptr[paged_kv.batch_size]; + + static_assert(num_stages_smem <= bdx); + uint32_t packed_page_iter_base = paged_kv.indptr[batch_idx] * paged_kv.page_size + chunk_start; +#pragma unroll + for (uint32_t j = 0; j < tile_size_per_bdx; ++j) { + uint32_t q, r; + paged_kv.page_size.divmod(packed_page_iter_base + ((j * bdz + tz) * bdy + ty) * bdx + tx, q, r); + kv_offset_smem[((j * bdz + tz) * bdy + ty) * bdx + tx] = + paged_kv.protective_get_kv_offset(q, kv_head_idx, r, 0, last_indptr); + } + block.sync(); + + size_t kv_offset[tile_size_per_bdx]; +#pragma unroll + for (uint32_t iter = 0; iter < num_stages_smem; ++iter) { +#pragma unroll + for (uint32_t j = 0; j < tile_size_per_bdx; ++j) { + kv_offset[j] = + kv_offset_smem[((iter * bdz + tz) * bdy + ty) * tile_size_per_bdx + j] + tx * vec_size; + } +#pragma unroll + for (uint32_t j = 0; j < tile_size_per_bdx; ++j) { + cp_async::pred_load( + k_smem + (((stage_idx * bdz + tz) * bdy + ty) * tile_size_per_bdx + j) * head_dim + + tx * vec_size, + paged_kv.k_data + kv_offset[j], + ((iter * bdz + tz) * bdy + ty) * tile_size_per_bdx + j < chunk_size); + } + cp_async::commit_group(); +#pragma unroll + for (uint32_t j = 0; j < tile_size_per_bdx; ++j) { + cp_async::pred_load( + v_smem + (((stage_idx * bdz + tz) * bdy + ty) * tile_size_per_bdx + j) * head_dim + + tx * vec_size, + paged_kv.v_data + kv_offset[j], + ((iter * bdz + tz) * bdy + ty) * tile_size_per_bdx + j < chunk_size); + } + cp_async::commit_group(); + stage_idx = (stage_idx + 1) % num_stages_smem; + } + + state_t st; + float s[bdy * tile_size_per_bdx]; + +#pragma unroll 2 + for (uint32_t iter = 0; iter < ceil_div(chunk_size, tile_size_per_bdx * bdy * bdz); ++iter) { + if ((iter + num_stages_smem) % bdx == 0) { +#pragma unroll + for (uint32_t j = 0; j < tile_size_per_bdx; ++j) { + uint32_t q, r; + paged_kv.page_size.divmod( + packed_page_iter_base + ((iter + num_stages_smem) * tile_size_per_bdx * bdy * bdz + + ((j * bdz + tz) * bdy + ty) * bdx + tx), + q, r); + kv_offset_smem[((j * bdz + tz) * bdy + ty) * bdx + tx] = + paged_kv.protective_get_kv_offset(q, kv_head_idx, r, 0, last_indptr); + } + } + // compute qk + cp_async::wait_group<2 * num_stages_smem - 1>(); + block.sync(); + compute_qk( + params, variant, batch_idx, + k_smem + (stage_idx * bdz + tz) * bdy * tile_size_per_bdx * head_dim, q_vec, freq, + (paged_kv.rope_pos_offset == nullptr ? 0 : paged_kv.rope_pos_offset[batch_idx]) + + chunk_start + iter * tile_size_per_bdx * bdy * bdz, + iter * tile_size_per_bdx * bdy * bdz, chunk_size, qo_head_idx, kv_head_idx, s, st, tx, ty, + tz); + block.sync(); + +#pragma unroll + for (uint32_t j = 0; j < tile_size_per_bdx; ++j) { + kv_offset[j] = kv_offset_smem[((((iter + num_stages_smem) % bdx) * bdz + tz) * bdy + ty) * + tile_size_per_bdx + + j] + + tx * vec_size; + } + + // load k tiles +#pragma unroll + for (uint32_t j = 0; j < tile_size_per_bdx; ++j) { + cp_async::pred_load( + k_smem + (((stage_idx * bdz + tz) * bdy + ty) * tile_size_per_bdx + j) * head_dim + + tx * vec_size, + paged_kv.k_data + kv_offset[j], + (((iter + num_stages_smem) * bdz + tz) * bdy + ty) * tile_size_per_bdx + j < chunk_size); + } + cp_async::commit_group(); + + // update m/d/o states + cp_async::wait_group<2 * num_stages_smem - 1>(); + block.sync(); + update_local_state( + v_smem + (stage_idx * bdz + tz) * bdy * tile_size_per_bdx * head_dim, s, stage_idx, st, tx); + block.sync(); + + // load v tiles +#pragma unroll + for (uint32_t j = 0; j < tile_size_per_bdx; ++j) { + cp_async::pred_load( + v_smem + (((stage_idx * bdz + tz) * bdy + ty) * tile_size_per_bdx + j) * head_dim + + tx * vec_size, + paged_kv.v_data + kv_offset[j], + (((iter + num_stages_smem) * bdz + tz) * bdy + ty) * tile_size_per_bdx + j < chunk_size); + } + cp_async::commit_group(); + stage_idx = (stage_idx + 1) % num_stages_smem; + } + cp_async::wait_group<0>(); + block.sync(); + + // sync local state of all warps inside a threadblock + sync_state(variant, st, reinterpret_cast(smem), smem_md, tx, ty, + tz); + if constexpr (variant.use_softmax) { + st.normalize(); + } + + if (tz == 0) { + st.o.cast_store(o + (bx * num_qo_heads + qo_head_idx) * head_dim + tx * vec_size); + // write lse + if (lse != nullptr) { + lse[bx * num_qo_heads + qo_head_idx] = st.get_lse(); + } + } +#if (__CUDACC_VER_MAJOR__ >= 12 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.launch_dependents;"); +#endif +} + +template +__global__ void BatchDecodeWithPagedKVCacheKernel(const __grid_constant__ Params params) { + extern __shared__ uint8_t smem[]; + BatchDecodeWithPagedKVCacheDevice(params, smem); +} + +/*! + * \brief Get the heuristic number of threads per threadblock + * \param group_size The number of qo heads that maps to the same kv head in GQA. + * \param sizeof_dtype The size (in terms of bytes) of the input data type + */ +constexpr uint32_t get_heuristic_num_threads(uint32_t group_size, uint32_t sizeof_dtype) { + if (group_size == 8U) { + if (sizeof_dtype == 1U) { + return 256U; // not enough registers for 512 threads + } else { + return 512U; + } + } else { + return 128U; + } +} + +/*! + * \brief FlashAttention decoding with kv-cache for a single request + * \tparam DTypeQ A template type indicates the query data type + * \tparam DTypeKV A template type indicates the key-value data type + * \tparam DTypeO A template type indicates the output data type + * \param q The query matrix, shape: [num_qo_heads, head_dim] + * \param k The key matrix in kv-cache, shape: [seq_len, num_kv_heads, head_dim] + * for NHD layout, [num_kv_heads, seq_len, head_dim] for HND layout + * \param v The value matrix in kv-cache, shape: [seq_len, num_kv_heads, + * head_dim] for NHD layout, [num_kv_heads, seq_len, head_dim] for HND layout + * \param o The output matrix, shape: [num_qo_heads, head_dim] + * \param tmp Used-allocated temporary buffer + * \param num_qo_heads A integer indicates the number of heads of query and output + * \param num_kv_heads A integer indicates the number of heads of key and value + * \param seq_len A integer indicates the sequence length + * \param head_dim A integer indicates the head dimension + * \param pos_encoding_mode The positional encoding mode + * \param rope_scale The scaling factor used in RoPE Interpolation + * \param rope_theta The theta used in RoPE + * \param stream The cuda stream to launch the kernel + * \return status Indicates whether CUDA calls are successful + */ +template +cudaError_t SingleDecodeWithKVCacheDispatched(Params params, typename Params::DTypeO* tmp, + cudaStream_t stream) { + using DTypeQ = typename Params::DTypeQ; + using DTypeKV = typename Params::DTypeKV; + using DTypeO = typename Params::DTypeO; + const uint32_t num_qo_heads = params.num_qo_heads; + const uint32_t num_kv_heads = params.num_kv_heads; + const uint32_t seq_len = params.kv_len; + + constexpr uint32_t vec_size = std::max(16UL / sizeof(DTypeKV), HEAD_DIM / 32UL); + constexpr uint32_t bdx = HEAD_DIM / vec_size; + auto compute_capacity = GetCudaComputeCapability(); + static_assert(bdx <= 32U); + DISPATCH_GQA_GROUP_SIZE(num_qo_heads / num_kv_heads, GROUP_SIZE, { + constexpr uint32_t bdy = GROUP_SIZE; + constexpr uint32_t num_threads = + std::max(get_heuristic_num_threads(GROUP_SIZE, sizeof(DTypeKV)), bdx * bdy); + constexpr uint32_t bdz = num_threads / (bdx * bdy); + constexpr uint32_t tile_size_per_bdx = GROUP_SIZE == 1 ? (sizeof(DTypeKV) == 1 ? 2U : 8U) : 1U; + DISPATCH_COMPUTE_CAP_DECODE_NUM_STAGES_SMEM(compute_capacity, NUM_STAGES_SMEM, { + const uint32_t smem_size = + 2U * NUM_STAGES_SMEM * bdy * tile_size_per_bdx * bdz * HEAD_DIM * sizeof(DTypeKV) + + 2U * bdy * bdz * sizeof(float); + auto kernel = + SingleDecodeWithKVCacheKernel; + FLASHINFER_CUDA_CALL( + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + + if (seq_len <= 256 || tmp == nullptr) { + // no need to use partition-kv kernel + dim3 nblks = dim3(1, num_kv_heads); + dim3 nthrs = dim3(bdx, bdy, bdz); + params.kv_chunk_size = seq_len; + void* args[] = {(void*)¶ms}; + FLASHINFER_CUDA_CALL( + cudaLaunchKernel((void*)kernel, nblks, nthrs, args, smem_size, stream)); + } else { + // use partition-kv kernel + int num_blocks_per_sm = 0; + int num_sm = 0; + int dev_id = 0; + FLASHINFER_CUDA_CALL(cudaGetDevice(&dev_id)); + FLASHINFER_CUDA_CALL( + cudaDeviceGetAttribute(&num_sm, cudaDevAttrMultiProcessorCount, dev_id)); + FLASHINFER_CUDA_CALL(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &num_blocks_per_sm, kernel, num_threads, smem_size)); + uint32_t max_grid_size = uint32_t(num_blocks_per_sm) * uint32_t(num_sm); + uint32_t max_num_kv_chunks = max_grid_size / num_kv_heads; + uint32_t kv_chunk_size = max(ceil_div(seq_len, max_num_kv_chunks), 256); + uint32_t num_chunks = ceil_div(seq_len, kv_chunk_size); + dim3 nblks = dim3(num_chunks, num_kv_heads); + if (nblks.x == 0 || nblks.y == 0) { + std::ostringstream err_msg; + err_msg << "Invalid kernel configuration: nblks=(" << nblks.x << "," << nblks.y << ")"; + FLASHINFER_ERROR(err_msg.str()); + } + dim3 nthrs = dim3(bdx, bdy, bdz); + float* tmp_lse = (float*)(tmp + num_chunks * num_qo_heads * HEAD_DIM); + auto o = params.o; + auto lse = params.lse; + params.o = tmp; + params.lse = tmp_lse; + params.kv_chunk_size = kv_chunk_size; + void* args[] = {(void*)¶ms}; + FLASHINFER_CUDA_CALL( + cudaLaunchKernel((void*)kernel, nblks, nthrs, args, smem_size, stream)); + if constexpr (AttentionVariant::use_softmax) { + FLASHINFER_CUDA_CALL( + MergeStates(tmp, tmp_lse, o, lse, num_chunks, 1, num_qo_heads, HEAD_DIM, stream)); + } else { + FLASHINFER_CUDA_CALL(AttentionSum(tmp, o, num_chunks, 1, num_qo_heads, HEAD_DIM, stream)); + } + } + }); + }); + return cudaSuccess; +} + +template +cudaError_t BatchDecodeWithPagedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v, + float* tmp_s, bool enable_pdl, + cudaStream_t stream) { + using DTypeQ = typename Params::DTypeQ; + using DTypeKV = typename Params::DTypeKV; + using DTypeO = typename Params::DTypeO; + using IdType = typename Params::IdType; + const uint32_t num_qo_heads = params.num_qo_heads; + const uint32_t num_kv_heads = params.paged_kv.num_heads; + const uint32_t padded_batch_size = params.padded_batch_size; + + constexpr uint32_t vec_size = std::max(16UL / sizeof(DTypeKV), HEAD_DIM / 32UL); + auto compute_capacity = GetCudaComputeCapability(); + constexpr uint32_t bdx = HEAD_DIM / vec_size; + static_assert(bdx <= 32); + DISPATCH_GQA_GROUP_SIZE(num_qo_heads / num_kv_heads, GROUP_SIZE, { + constexpr uint32_t bdy = GROUP_SIZE; + constexpr uint32_t num_threads = std::max(128U, bdx * bdy); + constexpr uint32_t bdz = num_threads / (bdx * bdy); + constexpr uint32_t tile_size_per_bdx = GROUP_SIZE == 1 ? (sizeof(DTypeKV) == 1 ? 2U : 4U) : 1U; + DISPATCH_COMPUTE_CAP_DECODE_NUM_STAGES_SMEM(compute_capacity, NUM_STAGES_SMEM, { + const uint32_t smem_size = + 2 * NUM_STAGES_SMEM * tile_size_per_bdx * bdy * bdz * HEAD_DIM * sizeof(DTypeKV) + + std::max(tile_size_per_bdx * num_threads * sizeof(DTypeKV*), + 2 * bdy * bdz * sizeof(float)); + auto kernel = + BatchDecodeWithPagedKVCacheKernel; + FLASHINFER_CUDA_CALL( + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + dim3 nblks(padded_batch_size, num_kv_heads); + dim3 nthrs(bdx, bdy, bdz); + + // PDL launch config + cudaLaunchAttribute attribute[1]; + cudaLaunchConfig_t config; + if (enable_pdl) { + attribute[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attribute[0].val.programmaticStreamSerializationAllowed = 1; + config.attrs = attribute; + config.numAttrs = 1; + config.gridDim = nblks; + config.blockDim = nthrs; + config.dynamicSmemBytes = smem_size; + config.stream = stream; + } + if (tmp_v == nullptr) { + // do not use partition-kv kernel + params.partition_kv = false; + + if (enable_pdl) { + FLASHINFER_CUDA_CALL(cudaLaunchKernelEx(&config, kernel, params)); + } else { + void* args[] = {(void*)¶ms}; + FLASHINFER_CUDA_CALL( + cudaLaunchKernel((void*)kernel, nblks, nthrs, args, smem_size, stream)); + } + } else { + // use partition-kv kernel + params.partition_kv = true; + auto o = params.o; + auto lse = params.lse; + params.o = tmp_v; + params.lse = tmp_s; + if (enable_pdl) { + FLASHINFER_CUDA_CALL(cudaLaunchKernelEx(&config, kernel, params)); + } else { + void* args[] = {(void*)¶ms}; + FLASHINFER_CUDA_CALL( + cudaLaunchKernel((void*)kernel, nblks, nthrs, args, smem_size, stream)); + } + if constexpr (AttentionVariant::use_softmax) { + FLASHINFER_CUDA_CALL(VariableLengthMergeStates( + tmp_v, tmp_s, params.o_indptr, o, lse, params.paged_kv.batch_size, nullptr, + num_qo_heads, HEAD_DIM, enable_pdl, stream)); + } else { + FLASHINFER_CUDA_CALL( + VariableLengthAttentionSum(tmp_v, params.o_indptr, o, params.paged_kv.batch_size, + nullptr, num_qo_heads, HEAD_DIM, enable_pdl, stream)); + } + } + }); + }); + return cudaSuccess; +} + +template +__device__ __forceinline__ void compute_qk_and_update_local_stat_mla( + const Params& params, AttentionVariant variant, const uint32_t batch_idx, const T* ckv_smem, + const vec_t& q_nope_vec, const T* kpe_smem, + const vec_t& q_pe_vec, const vec_t& freq, + uint32_t kv_idx_base, uint32_t iter_base, uint32_t iter_bound, state_t& st) { + uint32_t tx = threadIdx.x, tz = threadIdx.z; + constexpr uint32_t head_dim_ckv = bdx * vec_size_ckv; + constexpr uint32_t head_dim_kpe = bdx * vec_size_kpe; + float s[tile_size]; + float m_prev = st.m; +#pragma unroll + for (uint32_t j = 0; j < tile_size; ++j) { + vec_t ckv_vec; + ckv_vec.cast_load(ckv_smem + j * head_dim_ckv + tx * vec_size_ckv); + + vec_t kpe_vec; + kpe_vec.cast_load(kpe_smem + j * head_dim_kpe + tx * vec_size_kpe); + + s[j] = 0.f; +#pragma unroll + for (uint32_t i = 0; i < vec_size_ckv; ++i) { + s[j] += q_nope_vec[i] * ckv_vec[i]; + } +#pragma unroll + for (uint32_t i = 0; i < vec_size_kpe; ++i) { + s[j] += q_pe_vec[i] * kpe_vec[i]; + } + s[j] *= params.sm_scale; +#pragma unroll + for (uint32_t offset = bdx / 2; offset > 0; offset /= 2) { + s[j] += math::shfl_xor_sync(s[j], offset); + } + s[j] = (iter_base + tz * tile_size + j < iter_bound) ? s[j] : -math::inf; + st.m = max(st.m, s[j]); + } + + float o_scale = math::ptx_exp2(m_prev - st.m); + st.d *= o_scale; +#pragma unroll + for (uint32_t j = 0; j < tile_size; ++j) { + s[j] = math::ptx_exp2(s[j] - st.m); + st.d += s[j]; + } +#pragma unroll + for (uint32_t i = 0; i < vec_size_ckv; ++i) { + st.o[i] = st.o[i] * o_scale; + } + +#pragma unroll + for (uint32_t j = 0; j < tile_size; ++j) { + vec_t v_vec; + v_vec.cast_load(ckv_smem + j * head_dim_ckv + tx * vec_size_ckv); +#pragma unroll + for (uint32_t i = 0; i < vec_size_ckv; ++i) { + st.o[i] = st.o[i] + s[j] * v_vec[i]; + } + } +} + +template +__global__ void BatchDecodeWithPagedKVCacheKernelMLA(Params params) { + auto block = cg::this_thread_block(); + using DTypeQ = typename Params::DTypeQ; + using DTypeKV = typename Params::DTypeKV; + using DTypeO = typename Params::DTypeO; + using IdType = typename Params::IdType; + const DTypeQ* q_nope = params.q_nope; + const DTypeQ* q_pe = params.q_pe; + DTypeO* o = params.o; + float* lse = params.lse; + const auto& paged_kv = params.paged_kv; + const IdType* q_rope_offset = params.q_rope_offset; + const bool* block_valid_mask = params.block_valid_mask; + const uint32_t num_qo_heads = params.num_qo_heads; + const float rope_rcp_scale = params.rope_rcp_scale; + const float rope_rcp_theta = params.rope_rcp_theta; + const bool partition_kv = params.partition_kv; + params.sm_scale *= math::log2e; + + constexpr uint32_t head_dim_ckv = bdx * vec_size_ckv; + constexpr uint32_t head_dim_kpe = bdx * vec_size_kpe; + const uint32_t batch_idx = blockIdx.x; + const uint32_t tx = threadIdx.x, ty = threadIdx.y, tz = threadIdx.z; + const uint32_t t_offset = dim3_offset(bdy, bdx, tz, ty, tx); + + // NOTE(Zihao): when CUDAGraph is enabled, we will launch more blocks than + // the actual batch size, so we need to check if the current batch is valid + if (block_valid_mask && !block_valid_mask[batch_idx]) return; + const uint32_t mapped_batch_idx = params.request_indices[batch_idx]; + + const uint32_t orig_seq_len = paged_kv.get_length(mapped_batch_idx); + int32_t q_rope_offset_val = + q_rope_offset == nullptr ? (orig_seq_len - 1) : q_rope_offset[mapped_batch_idx]; + + const uint32_t kv_chunk_idx_in_orig_mapped_batch = params.kv_tile_indices[batch_idx]; + const uint32_t kv_chunk_size = *(params.kv_chunk_size_ptr); + const uint32_t cur_chunk_start = + partition_kv ? kv_chunk_idx_in_orig_mapped_batch * kv_chunk_size : 0; + const uint32_t cur_chunk_end = + partition_kv ? min((kv_chunk_idx_in_orig_mapped_batch + 1) * kv_chunk_size, orig_seq_len) + : orig_seq_len; + const uint32_t cur_chunk_len = cur_chunk_end - cur_chunk_start; + + uint32_t packed_page_iter_base = + paged_kv.indptr[mapped_batch_idx] * paged_kv.page_size + cur_chunk_start; + const IdType last_indptr = paged_kv.indptr[paged_kv.batch_size]; + + constexpr uint32_t kv_iter_len = bdy * bdz; + constexpr uint32_t compute_qk_tile = bdy; + + extern __attribute__((shared)) uint8_t smem[]; + DTypeKV* ckv_smem = (DTypeKV*)smem; + DTypeKV* kpe_smem = (DTypeKV*)((uint8_t*)ckv_smem + + num_stages_smem * kv_iter_len * head_dim_ckv * sizeof(DTypeKV)); + size_t* ckv_offset_smem = (size_t*)((uint8_t*)kpe_smem + num_stages_smem * kv_iter_len * + head_dim_kpe * sizeof(DTypeKV)); + size_t* kpe_offset_smem = (size_t*)((uint8_t*)ckv_offset_smem + bdx * bdy * bdz * sizeof(size_t)); + float* smem_md = (float*)ckv_offset_smem; + + AttentionVariant variant(params, batch_idx, smem); + + vec_t q_nope_vec[tile_size_qo_heads]; + vec_t q_pe_vec[tile_size_qo_heads]; + state_t st[tile_size_qo_heads]; + uint32_t qo_head_idx[tile_size_qo_heads]; + + vec_t freq; + +#pragma unroll + for (uint32_t i = 0; i < vec_size_kpe; ++i) { + freq[i] = rope_rcp_scale * __powf(rope_rcp_theta, float(2 * ((tx * vec_size_kpe + i) / 2)) / + float(head_dim_kpe)); + } +#if (__CUDACC_VER_MAJOR__ >= 12 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.wait;"); +#endif + // load q_nope and q_pe tile +#pragma unroll + for (int i = 0; i < tile_size_qo_heads; ++i) { + qo_head_idx[i] = dim3_offset(bdy, tile_size_qo_heads, blockIdx.y, threadIdx.y, i); + if (qo_head_idx[i] < num_qo_heads) { + q_nope_vec[i].cast_load(q_nope + + (mapped_batch_idx * num_qo_heads + qo_head_idx[i]) * head_dim_ckv + + tx * vec_size_ckv); + q_pe_vec[i].cast_load(q_pe + + (mapped_batch_idx * num_qo_heads + qo_head_idx[i]) * head_dim_kpe + + tx * vec_size_kpe); + } + } + + // init paged-cache read offset to be used + uint32_t q, r; + paged_kv.page_size.divmod(packed_page_iter_base + t_offset, q, r); + ckv_offset_smem[t_offset] = paged_kv.protective_get_offset_ckv(q, r, /*feat_idx*/ 0, last_indptr); + kpe_offset_smem[t_offset] = paged_kv.protective_get_offset_kpe(q, r, /*feat_idx*/ 0, last_indptr); + block.sync(); + + uint32_t stage_idx = 0; + constexpr uint32_t vec_bits = sizeof(DTypeKV) * vec_size_ckv * 8; + constexpr uint32_t tx_fold = vec_size_ckv / vec_size_kpe; + static_assert(num_stages_smem <= bdx); + size_t offset_bytes; + bool is_valid_range; +#pragma unroll + for (uint32_t iter = 0; iter < num_stages_smem; ++iter) { + is_valid_range = (iter * kv_iter_len + dim2_offset(bdy, tz, ty)) < cur_chunk_len; + + offset_bytes = ckv_offset_smem[dim3_offset(bdz, bdy, iter, tz, ty)] + tx * vec_size_ckv; + cp_async::pred_load( + ckv_smem + (stage_idx * kv_iter_len + dim2_offset(bdy, tz, ty)) * head_dim_ckv + + tx * vec_size_ckv, + paged_kv.ckv_data + offset_bytes, is_valid_range); + + offset_bytes = + kpe_offset_smem[dim3_offset(bdz, bdy, iter, tz, ty)] + tx / tx_fold * vec_size_ckv; + cp_async::pred_load( + kpe_smem + (stage_idx * kv_iter_len + dim2_offset(bdy, tz, ty)) * head_dim_kpe + + tx / tx_fold * vec_size_ckv, + paged_kv.kpe_data + offset_bytes, is_valid_range); + + cp_async::commit_group(); + stage_idx = (stage_idx + 1) % num_stages_smem; + } + +#pragma unroll + for (uint32_t iter = 0; iter < ceil_div(cur_chunk_len, kv_iter_len); ++iter) { + cp_async::wait_group<1 * num_stages_smem - 1>(); + block.sync(); + const int32_t kv_idx_base = + (paged_kv.rope_pos_offset == nullptr ? 0 : paged_kv.rope_pos_offset[mapped_batch_idx]) + + cur_chunk_start + iter * kv_iter_len; +#pragma unroll + for (int i = 0; i < tile_size_qo_heads; ++i) { + compute_qk_and_update_local_stat_mla( + params, variant, mapped_batch_idx, + ckv_smem + (stage_idx * kv_iter_len + tz * compute_qk_tile) * head_dim_ckv, q_nope_vec[i], + kpe_smem + (stage_idx * kv_iter_len + tz * compute_qk_tile) * head_dim_kpe, q_pe_vec[i], + freq, kv_idx_base, + /*iter_base*/ iter * kv_iter_len, /*iter_bound*/ cur_chunk_len, st[i]); + } + + if ((iter + num_stages_smem) % bdx == 0) { + uint32_t q, r; + paged_kv.page_size.divmod( + packed_page_iter_base + (iter + num_stages_smem) * kv_iter_len + t_offset, q, r); + ckv_offset_smem[t_offset] = + paged_kv.protective_get_offset_ckv(q, r, /*feat_idx*/ 0, last_indptr); + kpe_offset_smem[t_offset] = + paged_kv.protective_get_offset_kpe(q, r, /*feat_idx*/ 0, last_indptr); + } + block.sync(); + + is_valid_range = + ((iter + num_stages_smem) * kv_iter_len + dim2_offset(bdy, tz, ty)) < cur_chunk_len; + offset_bytes = ckv_offset_smem[dim3_offset(bdz, bdy, (iter + num_stages_smem) % bdx, tz, ty)] + + tx * vec_size_ckv; + cp_async::pred_load( + ckv_smem + (stage_idx * kv_iter_len + dim2_offset(bdy, tz, ty)) * head_dim_ckv + + tx * vec_size_ckv, + paged_kv.ckv_data + offset_bytes, is_valid_range); + + offset_bytes = kpe_offset_smem[dim3_offset(bdz, bdy, (iter + num_stages_smem) % bdx, tz, ty)] + + tx / tx_fold * vec_size_ckv; + cp_async::pred_load( + kpe_smem + (stage_idx * kv_iter_len + dim2_offset(bdy, tz, ty)) * head_dim_kpe + + tx / tx_fold * vec_size_ckv, + paged_kv.kpe_data + offset_bytes, is_valid_range); + cp_async::commit_group(); + + stage_idx = (stage_idx + 1) % num_stages_smem; + } + cp_async::wait_group<0>(); + block.sync(); + + if (bdz != 1) { +#pragma unroll + for (int i = 0; i < tile_size_qo_heads; ++i) { + if (qo_head_idx[i] < num_qo_heads) + sync_state(variant, st[i], (float*)smem, smem_md, tx, ty, tz); + } + } + + if (tz == 0) { +#pragma unroll + for (int i = 0; i < tile_size_qo_heads; ++i) { + if (qo_head_idx[i] < num_qo_heads) { + st[i].normalize(); + st[i].o.cast_store(o + (batch_idx * num_qo_heads + qo_head_idx[i]) * head_dim_ckv + + tx * vec_size_ckv); + + if (lse != nullptr) { + lse[batch_idx * num_qo_heads + qo_head_idx[i]] = st[i].get_lse(); + } + } + } + } +#if (__CUDACC_VER_MAJOR__ >= 12 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.launch_dependents;"); +#endif +} + +template +cudaError_t BatchDecodeWithPagedKVCacheDispatchedMLA(Params params, typename Params::DTypeO* tmp_v, + float* tmp_s, bool enable_pdl, + cudaStream_t stream) { + using DTypeQ = typename Params::DTypeQ; + using DTypeKV = typename Params::DTypeKV; + using DTypeO = typename Params::DTypeO; + using IdType = typename Params::IdType; + const uint32_t num_qo_heads = params.num_qo_heads; + const uint32_t padded_batch_size = params.padded_batch_size; + + constexpr uint32_t vec_size_ckv = std::max(16UL / sizeof(DTypeKV), HEAD_DIM_CKV / 32UL); + constexpr uint32_t bdx = HEAD_DIM_CKV / vec_size_ckv; + constexpr uint32_t vec_size_kpe = HEAD_DIM_KPE / bdx; + + constexpr uint32_t bdy = 8; + constexpr uint32_t tile_size_qo_heads = 2; + constexpr uint32_t qo_heads_per_block = bdy * tile_size_qo_heads; + constexpr uint32_t num_threads = std::max(128U, bdx * bdy); + constexpr uint32_t bdz = num_threads / (bdx * bdy); + const uint32_t gdy = ceil_div(num_qo_heads, qo_heads_per_block); + + auto compute_capacity = GetCudaComputeCapability(); + DISPATCH_COMPUTE_CAP_DECODE_NUM_STAGES_SMEM(compute_capacity, NUM_STAGES_SMEM, { + const uint32_t smem_size = + NUM_STAGES_SMEM * bdy * bdz * (HEAD_DIM_CKV + HEAD_DIM_KPE) * sizeof(DTypeKV) + + std::max(num_threads * sizeof(size_t) * 2, 2 * bdy * bdz * sizeof(float)); + + auto kernel = + BatchDecodeWithPagedKVCacheKernelMLA; + FLASHINFER_CUDA_CALL( + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + + dim3 nblks(padded_batch_size, gdy); + dim3 nthrs(bdx, bdy, bdz); + + // PDL launch config + cudaLaunchAttribute attribute[1]; + cudaLaunchConfig_t config; + if (enable_pdl) { + attribute[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attribute[0].val.programmaticStreamSerializationAllowed = 1; + config.attrs = attribute; + config.numAttrs = 1; + config.gridDim = nblks; + config.blockDim = nthrs; + config.dynamicSmemBytes = smem_size; + config.stream = stream; + } + + if (tmp_v == nullptr) { + // do not use partition-kv kernel + params.partition_kv = false; + if (enable_pdl) { + FLASHINFER_CUDA_CALL(cudaLaunchKernelEx(&config, kernel, params)); + } else { + void* args[] = {(void*)¶ms}; + FLASHINFER_CUDA_CALL( + cudaLaunchKernel((void*)kernel, nblks, nthrs, args, smem_size, stream)); + } + } else { + // use partition-kv kernel + params.partition_kv = true; + auto o = params.o; + auto lse = params.lse; + params.o = tmp_v; + params.lse = tmp_s; + if (enable_pdl) { + FLASHINFER_CUDA_CALL(cudaLaunchKernelEx(&config, kernel, params)); + } else { + void* args[] = {(void*)¶ms}; + FLASHINFER_CUDA_CALL( + cudaLaunchKernel((void*)kernel, nblks, nthrs, args, smem_size, stream)); + } + FLASHINFER_CUDA_CALL(VariableLengthMergeStates( + tmp_v, tmp_s, params.o_indptr, o, lse, params.paged_kv.batch_size, nullptr, num_qo_heads, + HEAD_DIM_CKV, enable_pdl, stream)); + } + }); + return cudaSuccess; +} + +} // namespace flashinfer + +#endif // FLASHINFER_DECODE_CUH_ diff --git a/include/flashinfer/attention/decode_mla_cute_sm80.cuh b/include/flashinfer/attention/decode_mla_cute_sm80.cuh new file mode 100644 index 0000000000000000000000000000000000000000..28f3a543cbae6d82b31d60448aa91dac3dd26edd --- /dev/null +++ b/include/flashinfer/attention/decode_mla_cute_sm80.cuh @@ -0,0 +1,558 @@ +/* + * Copyright (c) 2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef FLASHINFER_DECODE_CUTE_SM80_CUH_ +#define FLASHINFER_DECODE_CUTE_SM80_CUH_ +#include +#include +#include +#include +#include + +#include + +#include "../cp_async.cuh" +#include "../math.cuh" +#include "../pos_enc.cuh" +#include "../utils.cuh" +#include "../vec_dtypes.cuh" +#include "cascade.cuh" +#include "cute/tensor.hpp" +#include "state.cuh" + +namespace flashinfer { + +using namespace cute; + +namespace cg = cooperative_groups; +using cp_async::PrefetchMode; +using cp_async::SharedMemFillMode; + +template +std::tuple LaunchSpecForDecodeKernelMlaCuteSM80( + const uint32_t num_qo_heads) { + // fixme: below types and consts are duplicated from the ones from MLA decode kernel, we may + // refactor the duplication later + constexpr int k_smem_stages = 2; + constexpr int k_kv_tile_len = 8; + constexpr int k_warp_rows = 4; + constexpr int k_warp_cols = 2; + + using LayoutQo = + Layout, Int>, Stride, _1>>; + + using LayoutQnope = + Layout, Int>, Stride, _1>>; + using LayoutQpe = + Layout, Int>, Stride, _1>>; + + using LayoutAtt = + Layout, Int>, Stride, _1>>; + + using LayoutOScaleVec = Layout>>; + + using LayoutSwizzleAtomKV = decltype(composition( + Swizzle<3, 3, 3>{}, make_layout(make_shape(_8{}, _64{}), make_stride(_64{}, _1{})))); + using LayoutSwizzleQnope = decltype(tile_to_shape(LayoutSwizzleAtomKV{}, LayoutQnope{}.shape())); + using LayoutSwizzleCkv = decltype(tile_to_shape( + LayoutSwizzleAtomKV{}, Shape, Int, Int>{})); + using LayoutSwizzleQpe = decltype(tile_to_shape(LayoutSwizzleAtomKV{}, LayoutQpe{}.shape())); + using LayoutSwizzleKpe = decltype(tile_to_shape( + LayoutSwizzleAtomKV{}, Shape, Int, Int>{})); + + uint32_t smem_size = k_warp_rows * 32 * sizeof(size_t) * 2 + + (cosize(LayoutSwizzleQnope{}) + cosize(LayoutSwizzleQpe{}) + + cosize(LayoutSwizzleCkv{}) + cosize(LayoutSwizzleKpe{})) * + sizeof(DTypeKV) + + cosize(LayoutAtt{}) * sizeof(float) + + cosize(LayoutOScaleVec{}) * sizeof(float) * 2; + + const uint32_t gdy = ceil_div(num_qo_heads, QO_TILE_LEN); + + return {smem_size, gdy, k_warp_rows * k_warp_cols}; +} + +template +__global__ void BatchDecodeWithPagedKVCacheKernelMlaCuteSM80(Params params) { + auto block = cooperative_groups::this_thread_block(); + + static_assert(std::is_same::value, "DTypeQ is expeted to be fp16"); + static_assert(std::is_same::value, + "DTypeKV is expeted to be fp16"); + static_assert(std::is_same::value, "DTypeO is expeted to be fp16"); + + using IdType = typename Params::IdType; + using DTypeKV = half; + const DTypeKV* q_nope_ptr = params.q_nope; + const DTypeKV* q_pe_ptr = params.q_pe; + DTypeKV* output_ptr = params.o; + float* lse = params.lse; + const auto& paged_kv = params.paged_kv; + const bool* block_valid_mask = params.block_valid_mask; + const uint32_t num_qo_heads = params.num_qo_heads; + const bool partition_kv = params.partition_kv; + + const uint32_t batch_idx = blockIdx.x; + const uint32_t tx = threadIdx.x; + + // when CUDAGraph is enabled, we will launch more blocks than + // the actual batch size, so we need to check if the current batch is valid + if (block_valid_mask && !block_valid_mask[batch_idx]) return; + const uint32_t mapped_batch_idx = params.request_indices[batch_idx]; + + const uint32_t orig_seq_len = paged_kv.get_length(mapped_batch_idx); + + const uint32_t kv_chunk_idx_in_orig_mapped_batch = params.kv_tile_indices[batch_idx]; + const uint32_t kv_chunk_size = *(params.kv_chunk_size_ptr); + const uint32_t cur_chunk_start = + partition_kv ? kv_chunk_idx_in_orig_mapped_batch * kv_chunk_size : 0; + const uint32_t cur_chunk_end = + partition_kv ? min((kv_chunk_idx_in_orig_mapped_batch + 1) * kv_chunk_size, orig_seq_len) + : orig_seq_len; + const uint32_t cur_chunk_len = cur_chunk_end - cur_chunk_start; + + uint32_t packed_page_iter_base = + paged_kv.indptr[mapped_batch_idx] * paged_kv.page_size + cur_chunk_start; + const IdType last_indptr = paged_kv.indptr[paged_kv.batch_size]; + + const auto sm_scale = params.sm_scale * math::log2e; + + constexpr int k_smem_stages = 2; + constexpr int k_kv_tile_len = 8; + constexpr int k_warp_rows = 4; + constexpr int k_warp_cols = 2; + + using LayoutQo = + Layout, Int>, Stride, _1>>; + + using LayoutQnope = + Layout, Int>, Stride, _1>>; + using LayoutQpe = + Layout, Int>, Stride, _1>>; + + using LayoutAtt = + Layout, Int>, Stride, _1>>; + + using LayoutOScaleVec = Layout>>; + + using LayoutSwizzleAtomKV = decltype(composition( + Swizzle<3, 3, 3>{}, make_layout(make_shape(_8{}, _64{}), make_stride(_64{}, _1{})))); + using LayoutSwizzleQnope = decltype(tile_to_shape(LayoutSwizzleAtomKV{}, LayoutQnope{}.shape())); + using LayoutSwizzleCkv = decltype(tile_to_shape( + LayoutSwizzleAtomKV{}, Shape, Int, Int>{})); + + using LayoutSwizzleQpe = decltype(tile_to_shape(LayoutSwizzleAtomKV{}, LayoutQpe{}.shape())); + using LayoutSwizzleKpe = decltype(tile_to_shape( + LayoutSwizzleAtomKV{}, Shape, Int, Int>{})); + + const uint32_t q_head_idx_start = mapped_batch_idx * num_qo_heads + blockIdx.y * QO_TILE_LEN; + const uint32_t o_head_idx_start = batch_idx * num_qo_heads + blockIdx.y * QO_TILE_LEN; + + Tensor gmem_q_nope_chunk = + make_tensor(make_gmem_ptr(q_nope_ptr + q_head_idx_start * HEAD_DIM_CKV), LayoutQnope{}); + Tensor gmem_q_pe_chunk = + make_tensor(make_gmem_ptr(q_pe_ptr + q_head_idx_start * HEAD_DIM_KPE), LayoutQpe{}); + Tensor gmem_output_chunk = + make_tensor(make_gmem_ptr(output_ptr + o_head_idx_start * HEAD_DIM_CKV), LayoutQo{}); + + extern __shared__ char smem_data[]; + size_t* ckv_offset_smem = (size_t*)smem_data; + size_t* kpe_offset_smem = ckv_offset_smem + k_warp_rows * 32; + Tensor smem_q_nope = make_tensor(make_smem_ptr((DTypeKV*)(kpe_offset_smem + k_warp_rows * 32)), + LayoutSwizzleQnope{}); + Tensor smem_q_pe = make_tensor( + make_smem_ptr(smem_q_nope.data() + cute::cosize(LayoutSwizzleQnope{})), LayoutSwizzleQpe{}); + + Tensor smem_ckv_chunk = make_tensor( + make_smem_ptr(smem_q_pe.data() + cute::cosize(LayoutSwizzleQpe{})), LayoutSwizzleCkv{}); + Tensor smem_kpe_chunk = make_tensor( + make_smem_ptr(smem_ckv_chunk.data() + cute::cosize(LayoutSwizzleCkv{})), LayoutSwizzleKpe{}); + + Tensor smem_att = make_tensor( + make_smem_ptr((float*)(smem_kpe_chunk.data().ptr_ + cute::cosize(LayoutSwizzleKpe{}))), + LayoutAtt{}); + + float* ptr_o_scale = (float*)(smem_att.data().ptr_ + cute::cosize(LayoutAtt{})); + Tensor smem_o_scale = make_tensor(make_smem_ptr(ptr_o_scale), LayoutOScaleVec{}); + float* ptr_denom = ptr_o_scale + cute::cosize(LayoutOScaleVec{}); + Tensor smem_denom = make_tensor(make_smem_ptr(ptr_denom), LayoutOScaleVec{}); + + constexpr uint32_t k_thr_g2s_tile_m = k_kv_tile_len; // 8 + constexpr uint32_t k_thr_g2s_tile_k = k_warp_rows * 32 / k_thr_g2s_tile_m; // 16 + auto layout_thr_g2s_tile = + make_layout(make_shape(Int{}, Int{}), LayoutLeft{}); + const uint32_t thr_m_idx_within_tile = + tx % k_thr_g2s_tile_m; // it's also kv-idx for ckv and kpe sequence + const uint32_t thr_k_idx_within_tile = tx / k_thr_g2s_tile_m; + + // load q data to smem + Tensor gmem_q_nope_chunk_128bit = recast(gmem_q_nope_chunk); + Tensor gmem_q_nope_part_128bit = + local_partition(gmem_q_nope_chunk_128bit, layout_thr_g2s_tile, tx); + Tensor smem_q_nope_128bit = recast(smem_q_nope); + Tensor smem_q_nope_part_128bit = local_partition(smem_q_nope_128bit, layout_thr_g2s_tile, tx); + if (tx < k_warp_rows * 32) { +#pragma unroll + for (int n = 0; n < size<0>(gmem_q_nope_part_128bit); ++n) +#pragma unroll + for (int k = 0; k < size<1>(gmem_q_nope_part_128bit); ++k) { + smem_q_nope_part_128bit(n, k) = gmem_q_nope_part_128bit(n, k); + } + if (thr_k_idx_within_tile < (HEAD_DIM_KPE * sizeof(DTypeKV) / sizeof(cute::uint128_t))) { + Tensor gmem_q_pe_chunk_128bit = recast(gmem_q_pe_chunk); + Tensor gmem_q_pe_part_128bit = + local_partition(gmem_q_pe_chunk_128bit, layout_thr_g2s_tile, tx); + Tensor smem_q_pe_128bit = recast(smem_q_pe); + Tensor smem_q_pe_part_128bit = local_partition(smem_q_pe_128bit, layout_thr_g2s_tile, tx); + static_assert(size<1>(gmem_q_pe_part_128bit) == 1); +#pragma unroll + for (int n = 0; n < size<0>(gmem_q_pe_part_128bit); ++n) { + smem_q_pe_part_128bit(n, _0{}) = gmem_q_pe_part_128bit(n, _0{}); + } + } + } + block.sync(); + + // initialize variables needed by phase2 + Tensor smem_ckv_chunk_128bit = recast(smem_ckv_chunk); + Tensor smem_ckv_load_part_128bit = + local_partition(smem_ckv_chunk_128bit, layout_thr_g2s_tile, tx); + Tensor smem_kpe_chunk_128bit = recast(smem_kpe_chunk); + Tensor smem_kpe_load_part_128bit = + local_partition(smem_kpe_chunk_128bit, layout_thr_g2s_tile, tx); + + constexpr uint32_t k_mma_att_tile_k = 16; + using TiledMmaAtt = + decltype(make_tiled_mma(MMA_Atom>{}, + make_layout(Shape, _1, _1>{}, LayoutRight{}))); + TiledMmaAtt tiled_mma_att; + auto thr_mma = tiled_mma_att.get_slice(tx); + + Tensor smem_q_nope_local_tiles = local_tile( + smem_q_nope, make_tile(Int{}, Int{}), make_coord(_0{}, _)); + Tensor reg_q_nope_tile_part = thr_mma.partition_fragment_A(smem_q_nope_local_tiles(_, _, 0)); + + Tensor smem_ckv_local_tiles = + local_tile(smem_ckv_chunk, make_tile(Int{}, Int{}), + make_coord(_0{}, _)); + Tensor reg_ckv_tile_part = thr_mma.partition_fragment_B(smem_ckv_local_tiles(_, _, _0{}, _0{})); + + auto s2r_tiled_copy_a = make_tiled_copy_A(Copy_Atom{}, tiled_mma_att); + auto s2r_thr_copy_a = s2r_tiled_copy_a.get_slice(tx); + Tensor smem_q_nope_tiles_part = s2r_thr_copy_a.partition_S(smem_q_nope_local_tiles); + Tensor reg_q_nope_tile_part_view = s2r_thr_copy_a.retile_D(reg_q_nope_tile_part); + + auto s2r_tiled_copy_b = make_tiled_copy_B(Copy_Atom{}, tiled_mma_att); + auto s2r_thr_copy_b = s2r_tiled_copy_b.get_slice(tx); + + Tensor smem_ckv_tiles_part = s2r_thr_copy_b.partition_S(smem_ckv_local_tiles); + Tensor reg_ckv_tile_part_view = s2r_thr_copy_b.retile_D(reg_ckv_tile_part); + + Tensor smem_q_pe_local_tiles = local_tile( + smem_q_pe, make_tile(Int{}, Int{}), make_coord(_0{}, _)); + Tensor reg_q_pe_tile_part = thr_mma.partition_fragment_A(smem_q_pe_local_tiles(_, _, _0{})); + + Tensor smem_kpe_local_tiles = + local_tile(smem_kpe_chunk, make_tile(Int{}, Int{}), + make_coord(_0{}, _)); + Tensor reg_kpe_tile_part = thr_mma.partition_fragment_B(smem_kpe_local_tiles(_, _, _0{}, _0{})); + + Tensor smem_q_pe_tiles_part = s2r_thr_copy_a.partition_S(smem_q_pe_local_tiles); + Tensor reg_q_pe_tile_part_view = s2r_thr_copy_a.retile_D(reg_q_pe_tile_part); + + Tensor smem_kpe_tiles_part = s2r_thr_copy_b.partition_S(smem_kpe_local_tiles); + Tensor reg_kpe_tile_part_view = s2r_thr_copy_b.retile_D(reg_kpe_tile_part); + + Tensor smem_att_part_c = thr_mma.partition_C(smem_att); + Tensor reg_att_part_c = make_fragment_like(smem_att_part_c); + + using LayoutOScaleMat = Layout, Int>, Stride<_1, _0>>; + Tensor o_scale_broadcast_mat = make_tensor((ptr_o_scale), LayoutOScaleMat{}); + Tensor denom_broadcast_mat = make_tensor(make_smem_ptr(ptr_denom), LayoutOScaleMat{}); + + // initialize variables needed by phase3 + using TiledMmaOutput = decltype(make_tiled_mma( + MMA_Atom>{}, + make_layout(Shape, Int, _1>{}, LayoutRight{}))); + TiledMmaOutput tiled_mma_output; + auto thr_mma_output = tiled_mma_output.get_slice(tx); + + Tensor smem_att_part_a = thr_mma_output.partition_A(smem_att); + Tensor reg_att_part_a = + thr_mma_output.partition_fragment_A(make_tensor((DTypeKV*)0x0, LayoutAtt{})); + + auto layout_ckv_trans = + make_layout(make_shape(Int{}, Int{}, Int{}), + make_stride(Int{}, _1{}, Int{})); + auto layout_ckv_trans_cps = composition(smem_ckv_chunk.layout(), layout_ckv_trans); + Tensor smem_ckv_trans = make_tensor(smem_ckv_chunk.data(), layout_ckv_trans_cps); + + auto s2r_tiled_copy_b_ckv = + make_tiled_copy_B(Copy_Atom{}, tiled_mma_output); + auto s2r_thr_copy_b_ckv = s2r_tiled_copy_b_ckv.get_slice(tx); + Tensor smem_v_part = s2r_thr_copy_b_ckv.partition_S(smem_ckv_trans); + + auto layout_ckv_trans_no_stage = + make_layout(make_shape(Int{}, Int{}), + make_stride(Int{}, _1{})); + Tensor reg_v_part = + thr_mma_output.partition_fragment_B(make_tensor((DTypeKV*)0x0, layout_ckv_trans_no_stage)); + Tensor reg_v_part_view = s2r_thr_copy_b_ckv.retile_D(reg_v_part); + + Tensor gmem_output_chunk_part = thr_mma_output.partition_C(gmem_output_chunk); + // Tensor reg_output_part = make_fragment_like(gmem_output_chunk_part); + Tensor reg_output_part = + thr_mma_output.partition_fragment_C(make_tensor((float*)0x0, LayoutQo{})); + clear(reg_output_part); + + Tensor o_scale_mat_part = thr_mma_output.partition_C(o_scale_broadcast_mat); + Tensor denom_mat_part = thr_mma_output.partition_C(denom_broadcast_mat); + + // init paged-cache read offset to be used + uint32_t q, r; + if (tx < k_warp_rows * 32) { + paged_kv.page_size.divmod(packed_page_iter_base + tx, q, r); + ckv_offset_smem[tx] = paged_kv.protective_get_offset_ckv(q, r, /*feat_idx*/ 0, last_indptr); + kpe_offset_smem[tx] = paged_kv.protective_get_offset_kpe(q, r, /*feat_idx*/ 0, last_indptr); + } + block.sync(); + + uint32_t stage_idx = 0; + size_t offset_bytes; + bool is_valid_range; + if (tx < k_warp_rows * 32) { +#pragma unroll + for (uint32_t iter = 0; iter < k_smem_stages; ++iter) { + uint32_t kv_idx = iter * k_kv_tile_len + thr_m_idx_within_tile; + is_valid_range = kv_idx < cur_chunk_len; + + offset_bytes = ckv_offset_smem[kv_idx]; + static_assert(size<0>(smem_ckv_load_part_128bit) == 1); +#pragma unroll + for (int k = 0; k < size<1>(smem_ckv_load_part_128bit); ++k) { + cp_async::pred_load<128, cp_async::PrefetchMode::kPrefetch, + cp_async::SharedMemFillMode::kNoFill>( + &smem_ckv_load_part_128bit(_0{}, k, stage_idx), + (cute::uint128_t*)(paged_kv.ckv_data + offset_bytes) + k * k_thr_g2s_tile_k + + thr_k_idx_within_tile, + is_valid_range); + } + + offset_bytes = kpe_offset_smem[kv_idx]; + is_valid_range = + is_valid_range && + (thr_k_idx_within_tile < (HEAD_DIM_KPE * sizeof(DTypeKV) / sizeof(cute::uint128_t))); + static_assert(size<0>(smem_kpe_load_part_128bit) == 1 && + size<1>(smem_kpe_load_part_128bit) == 1); + cp_async::pred_load<128, cp_async::PrefetchMode::kPrefetch, + cp_async::SharedMemFillMode::kNoFill>( + &smem_kpe_load_part_128bit(_0{}, _0{}, stage_idx), + (cute::uint128_t*)(paged_kv.kpe_data + offset_bytes) + thr_k_idx_within_tile, + is_valid_range); + + cp_async::commit_group(); + stage_idx = (stage_idx + 1) % k_smem_stages; + } + } + + // start rolling update + float row_max = -flashinfer::math::inf; + float row_denom = 1.0; + for (uint32_t iter = 0; iter < ceil_div(cur_chunk_len, k_kv_tile_len); ++iter) { + if (tx < k_warp_rows * 32) { + cp_async::wait_group<1 * k_smem_stages - 1>(); + } + block.sync(); + + if (tx < k_warp_rows * 32) { + clear(reg_att_part_c); +#pragma unroll + for (int k_tile = 0; k_tile < size<3>(smem_q_nope_tiles_part); ++k_tile) { + cute::copy(s2r_tiled_copy_a, smem_q_nope_tiles_part(_, _, _, k_tile), + reg_q_nope_tile_part_view); + cute::copy(s2r_tiled_copy_b, smem_ckv_tiles_part(_, _, _, k_tile, stage_idx), + reg_ckv_tile_part_view); + cute::gemm(tiled_mma_att, reg_att_part_c, reg_q_nope_tile_part, reg_ckv_tile_part, + reg_att_part_c); + } +#pragma unroll + for (int k_tile = 0; k_tile < size<3>(smem_q_pe_tiles_part); ++k_tile) { + cute::copy(s2r_tiled_copy_a, smem_q_pe_tiles_part(_, _, _, k_tile), + reg_q_pe_tile_part_view); + cute::copy(s2r_tiled_copy_b, smem_kpe_tiles_part(_, _, _, k_tile, stage_idx), + reg_kpe_tile_part_view); + cute::gemm(tiled_mma_att, reg_att_part_c, reg_q_pe_tile_part, reg_kpe_tile_part, + reg_att_part_c); + } +#pragma unroll + for (int i = 0; i < cute::size(reg_att_part_c); ++i) { + reg_att_part_c(i) *= sm_scale; + } + cute::copy(reg_att_part_c, smem_att_part_c); + } + block.sync(); + + // Phase2 compute softmax + if (tx < QO_TILE_LEN) { + uint32_t valid_kv_len = cur_chunk_len - iter * k_kv_tile_len; + valid_kv_len = (valid_kv_len < k_kv_tile_len) ? valid_kv_len : k_kv_tile_len; + + float row_max_prev = row_max; +#pragma unroll + for (int i = 0; i < k_kv_tile_len; ++i) { + if (i >= valid_kv_len) smem_att(tx, i) = -flashinfer::math::inf; + row_max = max(row_max, smem_att(tx, i)); + } + + float row_o_scale = math::ptx_exp2(row_max_prev - row_max); + smem_o_scale(tx) = row_o_scale; + + row_denom *= row_o_scale; +#pragma unroll + for (int i = 0; i < k_kv_tile_len; ++i) { + smem_att(tx, i) = math::ptx_exp2(smem_att(tx, i) - row_max); + row_denom += smem_att(tx, i); + } + smem_denom(tx) = row_denom; + } + block.sync(); + + // Phase3 compute output + + // below code block is executed by all 8 warps + { +#pragma unroll + for (int i = 0; i < cute::size(reg_output_part); ++i) + reg_output_part(i) = reg_output_part(i) * o_scale_mat_part(i); + + cute::copy(smem_att_part_a, reg_att_part_a); + cute::copy(s2r_tiled_copy_b_ckv, smem_v_part(_, _, _, stage_idx), reg_v_part_view); + cute::gemm(tiled_mma_output, reg_output_part, reg_att_part_a, reg_v_part, reg_output_part); + } + + if (tx < k_warp_rows * 32) { + // refill offset_smem + constexpr uint32_t how_many__kv_tile_len__in__offset_smem = k_warp_rows * 32 / k_kv_tile_len; + if (((iter + k_smem_stages) % how_many__kv_tile_len__in__offset_smem) == 0) { + uint32_t q, r; + paged_kv.page_size.divmod( + packed_page_iter_base + (iter + k_smem_stages) * k_kv_tile_len + tx, q, r); + ckv_offset_smem[tx] = paged_kv.protective_get_offset_ckv(q, r, 0, last_indptr); + kpe_offset_smem[tx] = paged_kv.protective_get_offset_kpe(q, r, 0, last_indptr); + } + } + block.sync(); + + if (tx < k_warp_rows * 32) { + // commit next async copy task to pipeline + uint32_t kv_idx = (iter + k_smem_stages) * k_kv_tile_len + thr_m_idx_within_tile; + is_valid_range = kv_idx < cur_chunk_len; + + offset_bytes = ckv_offset_smem[kv_idx % (k_warp_rows * 32)]; + static_assert(size<0>(smem_ckv_load_part_128bit) == 1); +#pragma unroll + for (int k = 0; k < size<1>(smem_ckv_load_part_128bit); ++k) { + cp_async::pred_load<128, cp_async::PrefetchMode::kPrefetch, + cp_async::SharedMemFillMode::kNoFill>( + &smem_ckv_load_part_128bit(_0{}, k, stage_idx), + (cute::uint128_t*)(paged_kv.ckv_data + offset_bytes) + k * k_thr_g2s_tile_k + + thr_k_idx_within_tile, + is_valid_range); + } + + offset_bytes = kpe_offset_smem[kv_idx % (k_warp_rows * 32)]; + is_valid_range = + is_valid_range && + (thr_k_idx_within_tile < (HEAD_DIM_KPE * sizeof(DTypeKV) / sizeof(cute::uint128_t))); + static_assert(size<0>(smem_kpe_load_part_128bit) == 1 && + size<1>(smem_kpe_load_part_128bit) == 1); + cp_async::pred_load<128, cp_async::PrefetchMode::kPrefetch, + cp_async::SharedMemFillMode::kNoFill>( + &smem_kpe_load_part_128bit(_0{}, _0{}, stage_idx), + (cute::uint128_t*)(paged_kv.kpe_data + offset_bytes) + thr_k_idx_within_tile, + is_valid_range); + cp_async::commit_group(); + } + + stage_idx = (stage_idx + 1) % k_smem_stages; + } // end for kv tile iteration + if (tx < k_warp_rows * 32) { + cp_async::wait_group<0>(); + } + block.sync(); + + // final output phase +#pragma unroll + for (int i = 0; i < cute::size(reg_output_part); ++i) { + reg_output_part(i) = reg_output_part(i) / denom_mat_part(i); + } + + cute::copy(reg_output_part, gmem_output_chunk_part); + + if (lse != nullptr && tx < QO_TILE_LEN) { + lse[o_head_idx_start + tx] = row_max + math::ptx_log2(row_denom); + } +} + +template +cudaError_t BatchDecodeWithPagedKVCacheDispatchedMlaCuteSM80(Params params, + typename Params::DTypeO* tmp_v, + float* tmp_s, cudaStream_t stream) { + using DTypeQ = typename Params::DTypeQ; + using DTypeKV = typename Params::DTypeKV; + using DTypeO = typename Params::DTypeO; + using IdType = typename Params::IdType; + const uint32_t num_qo_heads = params.num_qo_heads; + const uint32_t padded_batch_size = params.padded_batch_size; + + auto [smem_size, gdy, k_warps] = + LaunchSpecForDecodeKernelMlaCuteSM80( + num_qo_heads); + auto kernel = + BatchDecodeWithPagedKVCacheKernelMlaCuteSM80; + + FLASHINFER_CUDA_CALL( + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + + if (tmp_v == nullptr) { + // do not use partition-kv kernel + dim3 nblks(padded_batch_size, gdy); + dim3 nthrs(k_warps * 32); + params.partition_kv = false; + void* args[] = {(void*)¶ms}; + FLASHINFER_CUDA_CALL(cudaLaunchKernel((void*)kernel, nblks, nthrs, args, smem_size, stream)); + } else { + // use partition-kv kernel + params.partition_kv = true; + auto o = params.o; + auto lse = params.lse; + params.o = tmp_v; + params.lse = tmp_s; + void* args[] = {(void*)¶ms}; + dim3 nblks(padded_batch_size, gdy); + dim3 nthrs(k_warps * 32); + FLASHINFER_CUDA_CALL(cudaLaunchKernel((void*)kernel, nblks, nthrs, args, smem_size, stream)); + FLASHINFER_CUDA_CALL(VariableLengthMergeStates(tmp_v, tmp_s, params.o_indptr, o, lse, + params.paged_kv.batch_size, nullptr, + num_qo_heads, HEAD_DIM_CKV, stream)); + } + + return cudaSuccess; +} + +} // namespace flashinfer + +#endif // FLASHINFER_DECODE_CUTE_SM80_CUH_ diff --git a/include/flashinfer/attention/default_decode_params.cuh b/include/flashinfer/attention/default_decode_params.cuh new file mode 100644 index 0000000000000000000000000000000000000000..d06e46338c4c2b21890a7cf71af18df6cd490af2 --- /dev/null +++ b/include/flashinfer/attention/default_decode_params.cuh @@ -0,0 +1,279 @@ +/* + * Copyright (c) 2024 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef FLASHINFER_DECODE_PARAMS_CUH_ +#define FLASHINFER_DECODE_PARAMS_CUH_ + +#include + +#include + +#include "../layout.cuh" +#include "../page.cuh" + +namespace flashinfer { + +template +struct SingleDecodeParams { + using DTypeQ = DTypeQ_; + using DTypeKV = DTypeKV_; + using DTypeO = DTypeO_; + using IdType = int32_t; + DTypeQ* q; + DTypeKV* k; + DTypeKV* v; + DTypeO* o; + float* lse; + float* maybe_alibi_slopes; + uint32_t kv_len; + uint32_t num_qo_heads; + uint32_t num_kv_heads; + uint32_t q_stride_n; + uint32_t q_stride_h; + uint32_t kv_stride_n; + uint32_t kv_stride_h; + int32_t window_left; + float logits_soft_cap; + float sm_scale; + float rope_rcp_scale; + float rope_rcp_theta; + uint32_t kv_chunk_size; + + __device__ __host__ SingleDecodeParams() + : q(nullptr), + k(nullptr), + v(nullptr), + o(nullptr), + lse(nullptr), + maybe_alibi_slopes(nullptr), + kv_len(0), + num_qo_heads(0), + num_kv_heads(0), + q_stride_n(0), + q_stride_h(0), + kv_stride_n(0), + kv_stride_h(0), + window_left(0), + logits_soft_cap(0.0f), + sm_scale(0.0f), + rope_rcp_scale(0.0f), + rope_rcp_theta(0.0f), + kv_chunk_size(0) {} + + __device__ __host__ SingleDecodeParams(DTypeQ* q, DTypeKV* k, DTypeKV* v, DTypeO* o, + float* maybe_alibi_slopes, uint32_t seq_len, + uint32_t num_qo_heads, uint32_t num_kv_heads, + QKVLayout kv_layout, uint32_t head_dim, + int32_t window_left, float logits_soft_cap, float sm_scale, + float rope_scale, float rope_theta) + : q(q), + k(k), + v(v), + o(o), + lse(nullptr), + maybe_alibi_slopes(maybe_alibi_slopes), + kv_len(seq_len), + num_qo_heads(num_qo_heads), + num_kv_heads(num_kv_heads), + q_stride_n(num_qo_heads * head_dim), + q_stride_h(head_dim), + kv_stride_n((kv_layout == QKVLayout::kNHD) ? num_kv_heads * head_dim : head_dim), + kv_stride_h((kv_layout == QKVLayout::kNHD) ? head_dim : seq_len * head_dim), + window_left(window_left), + logits_soft_cap(logits_soft_cap), + sm_scale(sm_scale), + rope_rcp_scale(1.f / rope_scale), + rope_rcp_theta(1.f / rope_theta), + kv_chunk_size(0) {} + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { return 1; } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return kv_len; + } +}; + +template +struct BatchDecodeParams { + using DTypeQ = DTypeQ_; + using DTypeKV = DTypeKV_; + using DTypeO = DTypeO_; + using IdType = IdType_; + + DTypeQ* q; + IdType* q_rope_offset; + paged_kv_t paged_kv; + DTypeO* o; + float* lse; + float* maybe_alibi_slopes; + uint32_t padded_batch_size; + uint32_t num_qo_heads; + IdType q_stride_n; + IdType q_stride_h; + int32_t window_left; + float logits_soft_cap; + float sm_scale; + float rope_rcp_scale; + float rope_rcp_theta; + + IdType* request_indices; + IdType* kv_tile_indices; + IdType* o_indptr; + IdType* kv_chunk_size_ptr; + bool* block_valid_mask; + bool partition_kv; + + __device__ __host__ BatchDecodeParams() + : q(nullptr), + q_rope_offset(nullptr), + paged_kv(), + o(nullptr), + lse(nullptr), + maybe_alibi_slopes(nullptr), + padded_batch_size(0), + num_qo_heads(0), + q_stride_n(0), + q_stride_h(0), + window_left(0), + logits_soft_cap(0.0f), + sm_scale(0.0f), + rope_rcp_scale(0.0f), + rope_rcp_theta(0.0f), + request_indices(nullptr), + kv_tile_indices(nullptr), + o_indptr(nullptr), + kv_chunk_size_ptr(nullptr), + block_valid_mask(nullptr), + partition_kv(false) {} + + __device__ __host__ BatchDecodeParams(DTypeQ* q, IdType* q_rope_offset, + paged_kv_t paged_kv, DTypeO* o, float* lse, + float* maybe_alibi_slopes, uint32_t num_qo_heads, + IdType q_stride_n, IdType q_stride_h, int32_t window_left, + float logits_soft_cap, float sm_scale, float rope_scale, + float rope_theta) + : q(q), + q_rope_offset(q_rope_offset), + paged_kv(paged_kv), + o(o), + lse(lse), + maybe_alibi_slopes(maybe_alibi_slopes), + padded_batch_size(0), + num_qo_heads(num_qo_heads), + q_stride_n(q_stride_n), + q_stride_h(q_stride_h), + window_left(window_left), + logits_soft_cap(logits_soft_cap), + sm_scale(sm_scale), + rope_rcp_scale(1.f / rope_scale), + rope_rcp_theta(1.f / rope_theta), + request_indices(nullptr), + kv_tile_indices(nullptr), + o_indptr(nullptr), + kv_chunk_size_ptr(nullptr), + block_valid_mask(nullptr), + partition_kv(false) {} + + __host__ __device__ __forceinline__ int32_t get_qo_len(int32_t batch_idx) const { return 1; } + + __host__ __device__ __forceinline__ int32_t get_kv_len(int32_t batch_idx) const { + return paged_kv.get_length(batch_idx); + } +}; + +template +struct BatchDecodeParamsMLA { + using DTypeQ = DTypeQ_; + using DTypeKV = DTypeKV_; + using DTypeO = DTypeO_; + using IdType = IdType_; + + DTypeQ* q_nope; + DTypeQ* q_pe; + DTypeO* o; + float* lse; + float sm_scale; + + IdType* q_rope_offset; + paged_kv_mla_t paged_kv; + uint32_t padded_batch_size; + uint32_t num_qo_heads; + int32_t window_left; + float logits_soft_cap; + float rope_rcp_scale; + float rope_rcp_theta; + + IdType* request_indices; + IdType* kv_tile_indices; + IdType* o_indptr; + IdType* kv_chunk_size_ptr; + bool* block_valid_mask; + bool partition_kv; + + __device__ __host__ BatchDecodeParamsMLA() + : q_nope(nullptr), + q_pe(nullptr), + o(nullptr), + lse(nullptr), + sm_scale(0.0f), + q_rope_offset(nullptr), + paged_kv(), + padded_batch_size(0), + num_qo_heads(0), + window_left(0), + logits_soft_cap(0.0f), + rope_rcp_scale(0.0f), + rope_rcp_theta(0.0f), + request_indices(nullptr), + kv_tile_indices(nullptr), + o_indptr(nullptr), + kv_chunk_size_ptr(nullptr), + block_valid_mask(nullptr), + partition_kv(false) {} + + __device__ __host__ BatchDecodeParamsMLA(DTypeQ* q_nope, DTypeQ* q_pe, IdType* q_rope_offset, + paged_kv_mla_t paged_kv, DTypeO* o, + float* lse, uint32_t num_qo_heads, int32_t window_left, + float logits_soft_cap, float sm_scale, float rope_scale, + float rope_theta) + : q_nope(q_nope), + q_pe(q_pe), + o(o), + lse(lse), + sm_scale(sm_scale), + q_rope_offset(q_rope_offset), + paged_kv(paged_kv), + padded_batch_size(0), + num_qo_heads(num_qo_heads), + window_left(window_left), + logits_soft_cap(logits_soft_cap), + rope_rcp_scale(1.f / rope_scale), + rope_rcp_theta(1.f / rope_theta), + request_indices(nullptr), + kv_tile_indices(nullptr), + o_indptr(nullptr), + kv_chunk_size_ptr(nullptr), + block_valid_mask(nullptr), + partition_kv(false) {} + + __host__ __device__ __forceinline__ int32_t get_qo_len(int32_t batch_idx) const { return 1; } + __host__ __device__ __forceinline__ int32_t get_kv_len(int32_t batch_idx) const { + return paged_kv.get_length(batch_idx); + } +}; + +} // namespace flashinfer + +#endif // FLASHINFER_DECODE_PARAMS_CUH_ diff --git a/include/flashinfer/attention/default_prefill_params.cuh b/include/flashinfer/attention/default_prefill_params.cuh new file mode 100644 index 0000000000000000000000000000000000000000..2e857fcc729678e71c37dc6557c02589acbde2f4 --- /dev/null +++ b/include/flashinfer/attention/default_prefill_params.cuh @@ -0,0 +1,412 @@ +/* + * Copyright (c) 2024 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef FLASHINFER_PREFILL_PARAMS_CUH_ +#define FLASHINFER_PREFILL_PARAMS_CUH_ + +#include + +#include +#include + +#include "../page.cuh" + +namespace flashinfer { + +template +struct SinglePrefillParams { + using DTypeQ = DTypeQ_; + using DTypeKV = DTypeKV_; + using DTypeO = DTypeO_; + using IdType = int32_t; + DTypeQ* q; + DTypeKV* k; + DTypeKV* v; + uint8_t* maybe_custom_mask; + DTypeO* o; + float* lse; + float* maybe_alibi_slopes; + uint_fastdiv group_size; + uint32_t qo_len; + uint32_t kv_len; + uint32_t num_qo_heads; + uint32_t num_kv_heads; + uint32_t q_stride_n; + uint32_t q_stride_h; + uint32_t k_stride_n; + uint32_t k_stride_h; + uint32_t v_stride_n; + uint32_t v_stride_h; + uint32_t head_dim; + int32_t window_left; + float logits_soft_cap; + float sm_scale; + float rope_rcp_scale; + float rope_rcp_theta; + + uint32_t partition_kv; + + __host__ SinglePrefillParams() + : q(nullptr), + k(nullptr), + v(nullptr), + maybe_custom_mask(nullptr), + o(nullptr), + lse(nullptr), + maybe_alibi_slopes(nullptr), + group_size(), + qo_len(0), + kv_len(0), + num_qo_heads(0), + num_kv_heads(0), + q_stride_n(0), + q_stride_h(0), + k_stride_n(0), + k_stride_h(0), + v_stride_n(0), + v_stride_h(0), + head_dim(0), + window_left(0), + logits_soft_cap(0.0f), + sm_scale(0.0f), + rope_rcp_scale(0.0f), + rope_rcp_theta(0.0f), + partition_kv(false) {} + + __host__ SinglePrefillParams(DTypeQ* q, DTypeKV* k, DTypeKV* v, uint8_t* maybe_custom_mask, + DTypeO* o, float* lse, float* maybe_alibi_slopes, + uint32_t num_qo_heads, uint32_t num_kv_heads, uint32_t qo_len, + uint32_t kv_len, uint32_t q_stride_n, uint32_t q_stride_h, + uint32_t kv_stride_n, uint32_t kv_stride_h, uint32_t head_dim, + int32_t window_left, float logits_soft_cap, float sm_scale, + float rope_scale, float rope_theta) + : q(q), + k(k), + v(v), + maybe_custom_mask(maybe_custom_mask), + o(o), + lse(lse), + maybe_alibi_slopes(maybe_alibi_slopes), + group_size(num_qo_heads / num_kv_heads), + num_qo_heads(num_qo_heads), + num_kv_heads(num_kv_heads), + qo_len(qo_len), + kv_len(kv_len), + q_stride_n(q_stride_n), + q_stride_h(q_stride_h), + k_stride_n(kv_stride_n), + k_stride_h(kv_stride_h), + v_stride_n(kv_stride_n), + v_stride_h(kv_stride_h), + head_dim(head_dim), + window_left(window_left), + logits_soft_cap(logits_soft_cap), + sm_scale(sm_scale), + rope_rcp_scale(1. / rope_scale), + rope_rcp_theta(1. / rope_theta), + partition_kv(false) {} + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { + return qo_len; + } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return kv_len; + } +}; + +template +struct BatchPrefillRaggedParams { + using DTypeQ = DTypeQ_; + using DTypeKV = DTypeKV_; + using DTypeO = DTypeO_; + using IdType = IdType_; + + DTypeQ* q; + DTypeKV* k; + DTypeKV* v; + uint8_t* maybe_custom_mask; + IdType* q_indptr; + IdType* kv_indptr; + IdType* maybe_mask_indptr; + IdType* maybe_q_rope_offset; // maybe_q_rope_offset is only used for fused-rope attention + IdType* maybe_k_rope_offset; // maybe_k_rope_offset is only used for fused-rope attention + DTypeO* o; + float* lse; + float* maybe_alibi_slopes; + uint_fastdiv group_size; + uint32_t num_qo_heads; + uint32_t num_kv_heads; + uint32_t q_stride_n; + uint32_t q_stride_h; + uint32_t k_stride_n; + uint32_t k_stride_h; + uint32_t v_stride_n; + uint32_t v_stride_h; + int32_t window_left; + float logits_soft_cap; + float sm_scale; + float rope_rcp_scale; + float rope_rcp_theta; + + IdType* request_indices; + IdType* qo_tile_indices; + IdType* kv_tile_indices; + IdType* merge_indptr; + IdType* o_indptr; + IdType* kv_chunk_size_ptr; + bool* block_valid_mask; + uint32_t max_total_num_rows; + uint32_t* total_num_rows; + uint32_t padded_batch_size; + bool partition_kv; + uint32_t* maybe_prefix_len_ptr; + uint16_t* maybe_token_pos_in_items_ptr; + uint32_t token_pos_in_items_len; + uint16_t* maybe_max_item_len_ptr; + + __host__ BatchPrefillRaggedParams() + : q(nullptr), + k(nullptr), + v(nullptr), + maybe_custom_mask(nullptr), + q_indptr(nullptr), + kv_indptr(nullptr), + maybe_mask_indptr(nullptr), + maybe_q_rope_offset(nullptr), + maybe_k_rope_offset(nullptr), + o(nullptr), + lse(nullptr), + maybe_alibi_slopes(nullptr), + group_size(), + num_qo_heads(0), + num_kv_heads(0), + q_stride_n(0), + q_stride_h(0), + k_stride_n(0), + k_stride_h(0), + v_stride_n(0), + v_stride_h(0), + window_left(0), + logits_soft_cap(0.0f), + sm_scale(0.0f), + rope_rcp_scale(0.0f), + rope_rcp_theta(0.0f), + request_indices(nullptr), + qo_tile_indices(nullptr), + kv_tile_indices(nullptr), + merge_indptr(nullptr), + o_indptr(nullptr), + kv_chunk_size_ptr(nullptr), + block_valid_mask(nullptr), + max_total_num_rows(0), + total_num_rows(nullptr), + padded_batch_size(0), + partition_kv(false), + maybe_prefix_len_ptr(nullptr), + maybe_token_pos_in_items_ptr(nullptr), + token_pos_in_items_len(0), + maybe_max_item_len_ptr(nullptr) {} + + __host__ BatchPrefillRaggedParams(DTypeQ* q, DTypeKV* k, DTypeKV* v, uint8_t* maybe_custom_mask, + IdType* q_indptr, IdType* kv_indptr, IdType* maybe_mask_indptr, + IdType* maybe_q_rope_offset, IdType* maybe_k_rope_offset, + DTypeO* o, float* lse, float* maybe_alibi_slopes, + uint32_t num_qo_heads, uint32_t num_kv_heads, + uint32_t q_stride_n, uint32_t q_stride_h, uint32_t kv_stride_n, + uint32_t kv_stride_h, int32_t window_left, + float logits_soft_cap, float sm_scale, float rope_scale, + float rope_theta) + : q(q), + k(k), + v(v), + maybe_custom_mask(maybe_custom_mask), + q_indptr(q_indptr), + kv_indptr(kv_indptr), + maybe_mask_indptr(maybe_mask_indptr), + maybe_q_rope_offset(maybe_q_rope_offset), + maybe_k_rope_offset(maybe_k_rope_offset), + o(o), + lse(lse), + maybe_alibi_slopes(maybe_alibi_slopes), + group_size(num_qo_heads / num_kv_heads), + num_qo_heads(num_qo_heads), + num_kv_heads(num_kv_heads), + q_stride_n(q_stride_n), + q_stride_h(q_stride_h), + k_stride_n(kv_stride_n), + k_stride_h(kv_stride_h), + v_stride_n(kv_stride_n), + v_stride_h(kv_stride_h), + window_left(window_left), + logits_soft_cap(logits_soft_cap), + sm_scale(sm_scale), + rope_rcp_scale(1.f / rope_scale), + rope_rcp_theta(1.f / rope_theta), + request_indices(nullptr), + qo_tile_indices(nullptr), + kv_tile_indices(nullptr), + merge_indptr(nullptr), + o_indptr(nullptr), + kv_chunk_size_ptr(nullptr), + block_valid_mask(nullptr), + max_total_num_rows(0), + total_num_rows(nullptr), + padded_batch_size(0), + partition_kv(false), + maybe_prefix_len_ptr(nullptr), + maybe_token_pos_in_items_ptr(nullptr), + token_pos_in_items_len(0), + maybe_max_item_len_ptr(nullptr) {} + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { + return q_indptr[batch_idx + 1] - q_indptr[batch_idx]; + } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return kv_indptr[batch_idx + 1] - kv_indptr[batch_idx]; + } +}; + +template +struct BatchPrefillPagedParams { + using DTypeQ = DTypeQ_; + using DTypeKV = DTypeKV_; + using DTypeO = DTypeO_; + using IdType = IdType_; + + DTypeQ* q; + paged_kv_t paged_kv; + uint8_t* maybe_custom_mask; + IdType* q_indptr; + IdType* maybe_mask_indptr; + IdType* maybe_q_rope_offset; // maybe_q_rope_offset is only used for fused-rope attention + DTypeO* o; + float* lse; + float* maybe_alibi_slopes; + uint_fastdiv group_size; + uint32_t num_qo_heads; + IdType q_stride_n; + IdType q_stride_h; + int32_t window_left; + float logits_soft_cap; + float sm_scale; + float rope_rcp_scale; + float rope_rcp_theta; + + IdType* request_indices; + IdType* qo_tile_indices; + IdType* kv_tile_indices; + IdType* merge_indptr; + IdType* o_indptr; + bool* block_valid_mask; + IdType* kv_chunk_size_ptr; + uint32_t max_total_num_rows; + uint32_t* total_num_rows; + uint32_t padded_batch_size; + bool partition_kv; + uint32_t* maybe_prefix_len_ptr; + uint16_t* maybe_token_pos_in_items_ptr; + uint32_t token_pos_in_items_len; + uint16_t* maybe_max_item_len_ptr; + + __host__ BatchPrefillPagedParams() + : q(nullptr), + paged_kv(), + maybe_custom_mask(nullptr), + q_indptr(nullptr), + maybe_mask_indptr(nullptr), + maybe_q_rope_offset(nullptr), + o(nullptr), + lse(nullptr), + maybe_alibi_slopes(nullptr), + group_size(), + num_qo_heads(0), + q_stride_n(0), + q_stride_h(0), + window_left(0), + logits_soft_cap(0.0f), + sm_scale(0.0f), + rope_rcp_scale(0.0f), + rope_rcp_theta(0.0f), + request_indices(nullptr), + qo_tile_indices(nullptr), + kv_tile_indices(nullptr), + merge_indptr(nullptr), + o_indptr(nullptr), + block_valid_mask(nullptr), + kv_chunk_size_ptr(nullptr), + max_total_num_rows(0), + total_num_rows(nullptr), + padded_batch_size(0), + partition_kv(false), + maybe_prefix_len_ptr(nullptr), + maybe_token_pos_in_items_ptr(nullptr), + token_pos_in_items_len(0), + maybe_max_item_len_ptr(nullptr) {} + + __host__ BatchPrefillPagedParams(DTypeQ* q, paged_kv_t paged_kv, + uint8_t* maybe_custom_mask, IdType* q_indptr, + IdType* maybe_mask_indptr, IdType* maybe_q_rope_offset, + DTypeO* o, float* lse, float* maybe_alibi_slopes, + uint32_t num_qo_heads, IdType q_stride_n, IdType q_stride_h, + int32_t window_left, float logits_soft_cap, float sm_scale, + float rope_scale, float rope_theta) + : q(q), + paged_kv(paged_kv), + maybe_custom_mask(maybe_custom_mask), + q_indptr(q_indptr), + maybe_mask_indptr(maybe_mask_indptr), + maybe_q_rope_offset(maybe_q_rope_offset), + o(o), + lse(lse), + maybe_alibi_slopes(maybe_alibi_slopes), + group_size(num_qo_heads / paged_kv.num_heads), + num_qo_heads(num_qo_heads), + q_stride_n(q_stride_n), + q_stride_h(q_stride_h), + window_left(window_left), + logits_soft_cap(logits_soft_cap), + sm_scale(sm_scale), + rope_rcp_scale(1.f / rope_scale), + rope_rcp_theta(1.f / rope_theta), + request_indices(nullptr), + qo_tile_indices(nullptr), + kv_tile_indices(nullptr), + merge_indptr(nullptr), + o_indptr(nullptr), + block_valid_mask(nullptr), + kv_chunk_size_ptr(nullptr), + max_total_num_rows(0), + total_num_rows(nullptr), + padded_batch_size(0), + partition_kv(false), + maybe_prefix_len_ptr(nullptr), + maybe_token_pos_in_items_ptr(nullptr), + token_pos_in_items_len(0), + maybe_max_item_len_ptr(nullptr) {} + + __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const { + return q_indptr[batch_idx + 1] - q_indptr[batch_idx]; + } + + __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const { + return paged_kv.get_length(batch_idx); + } +}; + +} // namespace flashinfer + +#endif // FLASHINFER_DECODE_PARAMS_CUH_ diff --git a/include/flashinfer/attention/heap.h b/include/flashinfer/attention/heap.h new file mode 100644 index 0000000000000000000000000000000000000000..484669fbe24221ffbaa1de5abc20b321930822c5 --- /dev/null +++ b/include/flashinfer/attention/heap.h @@ -0,0 +1,66 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef FLASHINFER_ATTENTION_HEAP_H +#define FLASHINFER_ATTENTION_HEAP_H + +#include +#include +#include +#include + +namespace flashinfer { + +/*! + * \brief Heap data structure for (index, value) pairs + * \note minimal element on top + */ +class MinHeap { + public: + // first: index, second: cost + using Element = std::pair; + + MinHeap(int capacity) : heap_(capacity) { + for (int i = 0; i < capacity; ++i) { + heap_[i] = std::make_pair(i, 0.f); + } + } + + void insert(const Element& element) { + heap_.push_back(element); + std::push_heap(heap_.begin(), heap_.end(), compare); + } + + Element pop() { + std::pop_heap(heap_.begin(), heap_.end(), compare); + Element minElement = heap_.back(); + heap_.pop_back(); + return minElement; + } + + std::vector getHeap() const { return heap_; } + + private: + // Custom comparator for the min-heap: compare based on 'val' in the pair + static bool compare(const Element& a, const Element& b) { + return a.second > b.second; // create a min-heap based on val + } + + std::vector heap_; +}; + +} // namespace flashinfer + +#endif // FLASHINFER_ATTENTION_HEAP_H diff --git a/include/flashinfer/attention/hopper.cuh b/include/flashinfer/attention/hopper.cuh new file mode 100644 index 0000000000000000000000000000000000000000..07cdd26d87de2e0905a3b8de530ac31d18709a28 --- /dev/null +++ b/include/flashinfer/attention/hopper.cuh @@ -0,0 +1,278 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef FLASHINFER_ATTENTION_HOPPER_HEADER_CUH_ +#define FLASHINFER_ATTENTION_HOPPER_HEADER_CUH_ +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +#include "../permuted_smem.cuh" + +namespace flashinfer { + +using namespace cute::SM90::GMMA; + +// using WGMMA_NN_64x32x16_F32BF16BF16_SS = +template +struct WGMMA_ASYNC_SS {}; + +template +struct WGMMA_ASYNC_RS {}; + +#define EXPAND_FRAG_ARGS_4(x) x[0], x[1], x[2], x[3] +#define EXPAND_FRAG_ARGS_8(x) x[0], x[1], x[2], x[3], x[4], x[5], x[6], x[7] +#define EXPAND_FRAG_ARGS_16(x) \ + x[0], x[1], x[2], x[3], x[4], x[5], x[6], x[7], x[8], x[9], x[10], x[11], x[12], x[13], x[14], \ + x[15] +#define EXPAND_FRAG_ARGS_32(x) \ + x[0], x[1], x[2], x[3], x[4], x[5], x[6], x[7], x[8], x[9], x[10], x[11], x[12], x[13], x[14], \ + x[15], x[16], x[17], x[18], x[19], x[20], x[21], x[22], x[23], x[24], x[25], x[26], x[27], \ + x[28], x[29], x[30], x[31] +#define EXPAND_FRAG_ARGS_64(x) \ + x[0], x[1], x[2], x[3], x[4], x[5], x[6], x[7], x[8], x[9], x[10], x[11], x[12], x[13], x[14], \ + x[15], x[16], x[17], x[18], x[19], x[20], x[21], x[22], x[23], x[24], x[25], x[26], x[27], \ + x[28], x[29], x[30], x[31], x[32], x[33], x[34], x[35], x[36], x[37], x[38], x[39], x[40], \ + x[41], x[42], x[43], x[44], x[45], x[46], x[47], x[48], x[49], x[50], x[51], x[52], x[53], \ + x[54], x[55], x[56], x[57], x[58], x[59], x[60], x[61], x[62], x[63] +#define EXPAND_FRAG_ARGS_128(x) \ + x[0], x[1], x[2], x[3], x[4], x[5], x[6], x[7], x[8], x[9], x[10], x[11], x[12], x[13], x[14], \ + x[15], x[16], x[17], x[18], x[19], x[20], x[21], x[22], x[23], x[24], x[25], x[26], x[27], \ + x[28], x[29], x[30], x[31], x[32], x[33], x[34], x[35], x[36], x[37], x[38], x[39], x[40], \ + x[41], x[42], x[43], x[44], x[45], x[46], x[47], x[48], x[49], x[50], x[51], x[52], x[53], \ + x[54], x[55], x[56], x[57], x[58], x[59], x[60], x[61], x[62], x[63], x[64], x[65], x[66], \ + x[67], x[68], x[69], x[70], x[71], x[72], x[73], x[74], x[75], x[76], x[77], x[78], x[79], \ + x[80], x[81], x[82], x[83], x[84], x[85], x[86], x[87], x[88], x[89], x[90], x[91], x[92], \ + x[93], x[94], x[95], x[96], x[97], x[98], x[99], x[100], x[101], x[102], x[103], x[104], \ + x[105], x[106], x[107], x[108], x[109], x[110], x[111], x[112], x[113], x[114], x[115], \ + x[116], x[117], x[118], x[119], x[120], x[121], x[122], x[123], x[124], x[125], x[126], \ + x[127] + +template +struct WGMMA_ASYNC_RS<__half, float, 64, 16, 16, TransposeA, TransposeB, scaleA, scaleB> { + template + static __device__ __forceinline__ void op(uint32_t* a_frag, uint64_t desc_b, float* d_frag) { + constexpr auto scale_d = init ? ScaleOut::Zero : ScaleOut::One; + MMA_64x16x16_F32F16F16_RS::fma( + EXPAND_FRAG_ARGS_4(a_frag), desc_b, EXPAND_FRAG_ARGS_8(d_frag), scale_d); + } +}; + +template +struct WGMMA_ASYNC_SS<__half, float, 64, 16, 16, TransposeA, TransposeB, scaleA, scaleB> { + template + static __device__ __forceinline__ void op(uint64_t desc_a, uint64_t desc_b, float* d_frag) { + constexpr auto scale_d = init ? ScaleOut::Zero : ScaleOut::One; + MMA_64x16x16_F32F16F16_SS::fma( + desc_a, desc_b, EXPAND_FRAG_ARGS_8(d_frag), scale_d); + } +}; + +template +struct WGMMA_ASYNC_RS<__half, float, 64, 32, 16, TransposeA, TransposeB, scaleA, scaleB> { + template + static __device__ __forceinline__ void op(uint32_t* a_frag, uint64_t desc_b, float* d_frag) { + constexpr auto scale_d = init ? ScaleOut::Zero : ScaleOut::One; + MMA_64x32x16_F32F16F16_RS::fma( + EXPAND_FRAG_ARGS_4(a_frag), desc_b, EXPAND_FRAG_ARGS_16(d_frag), scale_d); + } +}; + +template +struct WGMMA_ASYNC_SS<__half, float, 64, 32, 16, TransposeA, TransposeB, scaleA, scaleB> { + template + static __device__ __forceinline__ void op(uint64_t desc_a, uint64_t desc_b, float* d_frag) { + constexpr auto scale_d = init ? ScaleOut::Zero : ScaleOut::One; + MMA_64x32x16_F32F16F16_SS::fma( + desc_a, desc_b, EXPAND_FRAG_ARGS_16(d_frag), scale_d); + } +}; + +template +struct WGMMA_ASYNC_RS<__half, float, 64, 64, 16, TransposeA, TransposeB, scaleA, scaleB> { + template + static __device__ __forceinline__ void op(uint32_t* a_frag, uint64_t desc_b, float* d_frag) { + constexpr auto scale_d = init ? ScaleOut::Zero : ScaleOut::One; + MMA_64x64x16_F32F16F16_RS::fma( + EXPAND_FRAG_ARGS_4(a_frag), desc_b, EXPAND_FRAG_ARGS_32(d_frag), scale_d); + } +}; + +template +struct WGMMA_ASYNC_SS<__half, float, 64, 64, 16, TransposeA, TransposeB, scaleA, scaleB> { + template + static __device__ __forceinline__ void op(uint64_t desc_a, uint64_t desc_b, float* d_frag) { + constexpr auto scale_d = init ? ScaleOut::Zero : ScaleOut::One; + MMA_64x64x16_F32F16F16_SS::fma( + desc_a, desc_b, EXPAND_FRAG_ARGS_32(d_frag), scale_d); + } +}; + +template +struct WGMMA_ASYNC_RS<__half, float, 64, 128, 16, TransposeA, TransposeB, scaleA, scaleB> { + template + static __device__ __forceinline__ void op(uint32_t* a_frag, uint64_t desc_b, float* d_frag) { + constexpr auto scale_d = init ? ScaleOut::Zero : ScaleOut::One; + MMA_64x128x16_F32F16F16_RS::fma( + EXPAND_FRAG_ARGS_4(a_frag), desc_b, EXPAND_FRAG_ARGS_64(d_frag), scale_d); + } +}; + +template +struct WGMMA_ASYNC_SS<__half, float, 64, 128, 16, TransposeA, TransposeB, scaleA, scaleB> { + template + static __device__ __forceinline__ void op(uint64_t desc_a, uint64_t desc_b, float* d_frag) { + constexpr auto scale_d = init ? ScaleOut::Zero : ScaleOut::One; + MMA_64x128x16_F32F16F16_SS::fma( + desc_a, desc_b, EXPAND_FRAG_ARGS_64(d_frag), scale_d); + } +}; + +template +struct WGMMA_ASYNC_RS<__half, float, 64, 256, 16, TransposeA, TransposeB, scaleA, scaleB> { + template + static __device__ __forceinline__ void op(uint32_t* a_frag, uint64_t desc_b, float* d_frag) { + constexpr auto scale_d = init ? ScaleOut::Zero : ScaleOut::One; + MMA_64x256x16_F32F16F16_RS::fma( + EXPAND_FRAG_ARGS_4(a_frag), desc_b, EXPAND_FRAG_ARGS_128(d_frag), scale_d); + } +}; + +template +struct WGMMA_ASYNC_SS<__half, float, 64, 256, 16, TransposeA, TransposeB, scaleA, scaleB> { + template + static __device__ __forceinline__ void op(uint64_t desc_a, uint64_t desc_b, float* d_frag) { + constexpr auto scale_d = init ? ScaleOut::Zero : ScaleOut::One; + MMA_64x256x16_F32F16F16_SS::fma( + desc_a, desc_b, EXPAND_FRAG_ARGS_128(d_frag), scale_d); + } +}; + +template +struct WGMMA_ASYNC_SS<__nv_bfloat16, float, 64, 16, 16, TransposeA, TransposeB, scaleA, scaleB> { + template + static __device__ __forceinline__ void op(uint64_t desc_a, uint64_t desc_b, float* d_frag) { + constexpr auto scale_d = init ? ScaleOut::Zero : ScaleOut::One; + MMA_64x16x16_F32BF16BF16_SS::fma( + desc_a, desc_b, EXPAND_FRAG_ARGS_8(d_frag), scale_d); + } +}; + +template +struct WGMMA_ASYNC_SS<__nv_bfloat16, float, 64, 32, 16, TransposeA, TransposeB, scaleA, scaleB> { + template + static __device__ __forceinline__ void op(uint64_t desc_a, uint64_t desc_b, float* d_frag) { + constexpr auto scale_d = init ? ScaleOut::Zero : ScaleOut::One; + MMA_64x32x16_F32BF16BF16_SS::fma( + desc_a, desc_b, EXPAND_FRAG_ARGS_16(d_frag), scale_d); + } +}; + +template +struct WGMMA_ASYNC_SS<__nv_bfloat16, float, 64, 64, 16, TransposeA, TransposeB, scaleA, scaleB> { + template + static __device__ __forceinline__ void op(uint64_t desc_a, uint64_t desc_b, float* d_frag) { + constexpr auto scale_d = init ? ScaleOut::Zero : ScaleOut::One; + MMA_64x64x16_F32BF16BF16_SS::fma( + desc_a, desc_b, EXPAND_FRAG_ARGS_32(d_frag), scale_d); + } +}; + +template +struct WGMMA_ASYNC_SS<__nv_bfloat16, float, 64, 128, 16, TransposeA, TransposeB, scaleA, scaleB> { + template + static __device__ __forceinline__ void op(uint64_t desc_a, uint64_t desc_b, float* d_frag) { + constexpr auto scale_d = init ? ScaleOut::Zero : ScaleOut::One; + MMA_64x128x16_F32BF16BF16_SS::fma( + desc_a, desc_b, EXPAND_FRAG_ARGS_64(d_frag), scale_d); + } +}; + +template +struct WGMMA_ASYNC_SS<__nv_bfloat16, float, 64, 256, 16, TransposeA, TransposeB, scaleA, scaleB> { + template + static __device__ __forceinline__ void op(uint64_t desc_a, uint64_t desc_b, float* d_frag) { + constexpr auto scale_d = init ? ScaleOut::Zero : ScaleOut::One; + MMA_64x256x16_F32BF16BF16_SS::fma( + desc_a, desc_b, EXPAND_FRAG_ARGS_128(d_frag), scale_d); + } +}; + +using Swizzle128B = cute::Swizzle<3, 4, 3>; +using Swizzle64B = cute::Swizzle<2, 4, 3>; +using Swizzle32B = cute::Swizzle<1, 4, 3>; + +template +__device__ __forceinline__ uint32_t get_swizzle_offset(uint32_t i, uint32_t j) { + constexpr uint32_t M = 8; + if constexpr (swizzle_mode == SwizzleMode::k128B) { + constexpr uint32_t N = 8; + return Swizzle128B{}(((i / M) * M * stride + ((j / N) * M + i % M) * N + (j % N)) << 4) >> 4; + } else { + constexpr uint32_t N = 4; + return Swizzle64B{}(((i / M) * M * stride + ((j / N) * M + i % M) * N + (j % N)) << 4) >> 4; + } +} + +__device__ __forceinline__ uint64_t matrix_descriptor_encode(uint64_t x) { + return (((x) & 0x3FFFF) >> 0x4); +} + +template +__device__ uint64_t make_smem_desc(T* ptr) { + uint32_t addr = static_cast(__cvta_generic_to_shared(ptr)); + uint64_t desc = 0x0000000000000000; + desc |= matrix_descriptor_encode(addr); + // leading byte offset + desc |= matrix_descriptor_encode(leading_byte_offset) << 16; + // stride byte offset + desc |= matrix_descriptor_encode(stride_byte_offset) << 32; + desc |= ((swizzle_mode == SwizzleMode::k128B) ? 1llu + : (swizzle_mode == SwizzleMode::k64B) ? 2llu + : 3llu) + << 62; + return desc; +} + +__device__ __forceinline__ void warpgroup_arrive() { cute::warpgroup_arrive(); } + +template +__device__ __forceinline__ void warpgroup_wait() { + cute::warpgroup_wait(); +} + +__device__ __forceinline__ void warpgroup_commit_batch() { cute::warpgroup_commit_batch(); } + +template +__device__ __forceinline__ void warpgroup_fence_frag(float* frag) { +#pragma unroll + for (uint32_t i = 0; i < size; ++i) { + cute::warpgroup_fence_operand(frag[i]); + } +} + +}; // namespace flashinfer + +#endif // FLASHINFER_ATTENTION_HOPPER_HEADER_CUH_ diff --git a/include/flashinfer/attention/hopper/attention_updater.cuh b/include/flashinfer/attention/hopper/attention_updater.cuh new file mode 100644 index 0000000000000000000000000000000000000000..a9327ff7dd0721c5775a4104a42faf1a763cfbb2 --- /dev/null +++ b/include/flashinfer/attention/hopper/attention_updater.cuh @@ -0,0 +1,258 @@ +/* + * Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri + * Dao. Licensed under the BSD 3-Clause. + * + * Modified by the FlashInfer team. + */ +#ifndef FLASHINFER_ATTENTION_HOPPER_ATTENTION_UPDATER_CUH_ +#define FLASHINFER_ATTENTION_HOPPER_ATTENTION_UPDATER_CUH_ + +#include +#include + +#include "flashinfer/attention/hopper/utils.cuh" + +namespace flashinfer { + +using namespace cute; + +template +struct MaxOp { + __device__ __forceinline__ T operator()(T const& x, T const& y) { return x > y ? x : y; } +}; + +template <> +struct MaxOp { + // This is slightly faster + __device__ __forceinline__ float operator()(float const& x, float const& y) { return max(x, y); } +}; + +template +struct SumOp { + __device__ __forceinline__ T operator()(T const& x, T const& y) { return x + y; } +}; + +template +struct Allreduce { + static_assert(THREADS == 32 || THREADS == 16 || THREADS == 8 || THREADS == 4); + template + static __device__ __forceinline__ T run(T x, Operator& op) { + constexpr int OFFSET = THREADS / 2; + x = op(x, __shfl_xor_sync(uint32_t(-1), x, OFFSET)); + return Allreduce::run(x, op); + } +}; + +template <> +struct Allreduce<2> { + template + static __device__ __forceinline__ T run(T x, Operator& op) { + x = op(x, __shfl_xor_sync(uint32_t(-1), x, 1)); + return x; + } +}; + +template +__device__ __forceinline__ void thread_reduce_(Tensor const& tensor, + Tensor& summary, Operator& op) { + static_assert(Layout0::rank == 2, "Only support 2D Tensor"); + static_assert(Layout1::rank == 1, "Only support 1D Tensor"); + CUTE_STATIC_ASSERT_V(size<0>(summary) == size<0>(tensor)); +#pragma unroll + for (int mi = 0; mi < size<0>(tensor); mi++) { + summary(mi) = init ? tensor(mi, 0) : op(summary(mi), tensor(mi, 0)); +#pragma unroll + for (int ni = 1; ni < size<1>(tensor); ni++) { + summary(mi) = op(summary(mi), tensor(mi, ni)); + } + } +} + +template +__device__ __forceinline__ void quad_allreduce_(Tensor& dst, + Tensor& src, Operator& op) { + CUTE_STATIC_ASSERT_V(size(dst) == size(src)); +#pragma unroll + for (int i = 0; i < size(dst); i++) { + dst(i) = Allreduce<4>::run(src(i), op); + } +} + +template +__device__ __forceinline__ void reduce_(Tensor const& tensor, + Tensor& summary, Operator& op) { + thread_reduce_(tensor, summary, op); + quad_allreduce_(summary, summary, op); +} + +template +__device__ __forceinline__ void reduce_max(Tensor const& tensor, + Tensor& max) { + MaxOp max_op; + reduce_(tensor, max, max_op); +} + +template +__device__ __forceinline__ void reduce_sum(Tensor const& tensor, + Tensor& sum) { + SumOp sum_op; + thread_reduce_(tensor, sum, sum_op); + if constexpr (warp_reduce) { + quad_allreduce_(sum, sum, sum_op); + } +} + +template +__forceinline__ __device__ void apply_exp2(Tensor& tensor, + Tensor const& max) { + static_assert(Layout0::rank == 2, "Only support 2D Tensor"); + static_assert(Layout1::rank == 1, "Only support 1D Tensor"); + CUTE_STATIC_ASSERT_V(size<0>(max) == size<0>(tensor)); +#pragma unroll + for (int mi = 0; mi < size<0>(tensor); ++mi) { + auto row_max = max(mi); +#pragma unroll + for (int ni = 0; ni < size<1>(tensor); ++ni) { + tensor(mi, ni) = exp2f(tensor(mi, ni) - row_max); + } + } +} + +template +__forceinline__ __device__ void scale_apply_exp2(Tensor& tensor, + Tensor const& max, + const float scale) { + static_assert(Layout0::rank == 2, "Only support 2D Tensor"); + static_assert(Layout1::rank == 1, "Only support 1D Tensor"); + CUTE_STATIC_ASSERT_V(size<0>(max) == size<0>(tensor)); +#pragma unroll + for (int mi = 0; mi < size<0>(tensor); ++mi) { + auto row_max = max(mi); +#pragma unroll + for (int ni = 0; ni < size<1>(tensor); ++ni) { + // row_max * scale is a constant for each row, so we can use fma here + tensor(mi, ni) = exp2f(tensor(mi, ni) * scale - row_max * scale); + } + } +} + +template +struct DefaultUpdater { + using TensorT = decltype(make_tensor(Shape>{})); + constexpr static float fill_value = 0.f; + CUTLASS_DEVICE DefaultUpdater() {}; + + __forceinline__ __device__ TensorT get_lse() { return TensorT(); } + + template + __forceinline__ __device__ void update(Tensor0& acc_s) { + // NOTE(Zihao): nothing to do here + }; + + template + __forceinline__ __device__ void finalize(Tensor1& acc_s, float pv_scale = 1.f) { + // NOTE(Zihao): nothing to do here + }; + + template + __forceinline__ __device__ void rescale_o(Tensor1& acc_o) { + // NOTE(Zihao): nothing to do here + }; +}; + +template +struct OnlineSoftmax { + constexpr static float fill_value = -math::inf; + using TensorT = decltype(make_tensor(Shape>{})); + TensorT row_max, row_sum, scores_scale; + float sm_scale_log2; + + CUTLASS_DEVICE OnlineSoftmax(float sm_scale_log2) : sm_scale_log2(sm_scale_log2) { + clear(scores_scale); + }; + + __forceinline__ __device__ TensorT get_lse() const { return row_sum; } + + template + __forceinline__ __device__ void update(Tensor0& acc_s) { + // Reshape acc_s from ((2, 2, V), MMA_M, MMA_N) to (nrow=(2, MMA_M), ncol=(2, V, MMA_N)) + Tensor scores = make_tensor(acc_s.data(), convert_layout_acc_rowcol(acc_s.layout())); + + static_assert(decltype(size<0>(scores))::value == NUM_ROWS_PER_THREAD); + if constexpr (init) { + reduce_max(scores, row_max); + if constexpr (WITH_SCALE) { + scale_apply_exp2(scores, row_max, sm_scale_log2); + } else { + apply_exp2(scores, row_max); + } + reduce_sum(scores, row_sum); + } else { + // update row_max + Tensor scores_max_prev = make_fragment_like(row_max); + cute::copy(row_max, scores_max_prev); + reduce_max(scores, row_max); + // update scores_scale and scale row_sum +#pragma unroll + for (int mi = 0; mi < size(row_max); ++mi) { + float scores_max_cur = row_max(mi); + if constexpr (WITH_SCALE) { + scores_scale(mi) = exp2f((scores_max_prev(mi) - scores_max_cur) * sm_scale_log2); + } else { + scores_scale(mi) = exp2f(scores_max_prev(mi) - scores_max_cur); + } + row_sum(mi) *= scores_scale(mi); + } + // perform exp2 on scores + if constexpr (WITH_SCALE) { + scale_apply_exp2(scores, row_max, sm_scale_log2); + } else { + apply_exp2(scores, row_max); + } + // update row_sum + reduce_sum(scores, row_sum); + } + }; + + template + __forceinline__ __device__ void finalize(Tensor0& acc_s, float pv_scale = 1.f) { + // Reshape acc_s from ((2, 2, V), MMA_M, MMA_N) to (nrow=(2, MMA_M), ncol=(2, V, MMA_N)) + // Note (Yilong): use pv_scale to dequantize the output + Tensor scores = make_tensor(acc_s.data(), convert_layout_acc_rowcol(acc_s.layout())); + static_assert(decltype(size<0>(scores))::value == NUM_ROWS_PER_THREAD); + SumOp sum_op; + quad_allreduce_(row_sum, row_sum, sum_op); +#pragma unroll + for (int mi = 0; mi < size(row_max); ++mi) { + float sum = row_sum(mi); + float inv_sum = pv_scale / sum; + scores_scale(mi) = inv_sum; + if constexpr (WITH_SCALE) { + row_sum(mi) = row_max(mi) * sm_scale_log2 + math::ptx_log2(sum); + } else { + row_sum(mi) = row_max(mi) + math::ptx_log2(sum); + } + } + }; + + template + __forceinline__ __device__ void rescale_o(Tensor1& acc_o) { + // Reshape acc_o from (MMA=4, MMA_M, MMA_K) to (nrow=(2, MMA_M), ncol=(2, MMA_K)) + Tensor acc_o_rowcol = make_tensor(acc_o.data(), convert_layout_acc_rowcol(acc_o.layout())); + static_assert(decltype(size<0>(acc_o_rowcol))::value == NUM_ROWS_PER_THREAD); +#pragma unroll + for (int mi = 0; mi < size(row_max); ++mi) { +#pragma unroll + for (int ni = 0; ni < size<1>(acc_o_rowcol); ++ni) { + acc_o_rowcol(mi, ni) *= scores_scale(mi); + } + } + }; +}; + +} // namespace flashinfer + +#endif // FLASHINFER_ATTENTION_HOPPER_ATTENTION_UPDATER_CUH_ diff --git a/include/flashinfer/attention/hopper/block_sparse_gather.cuh b/include/flashinfer/attention/hopper/block_sparse_gather.cuh new file mode 100644 index 0000000000000000000000000000000000000000..7082f1da1fb85531b8b5a57cdaa48ab0210cca80 --- /dev/null +++ b/include/flashinfer/attention/hopper/block_sparse_gather.cuh @@ -0,0 +1,196 @@ +/* + * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + * Modified by the FlashInfer team. + */ +#ifndef FLASHINFER_ATTENTION_HOPPER_BLOCK_SPARSE_GATHER_CUH +#define FLASHINFER_ATTENTION_HOPPER_BLOCK_SPARSE_GATHER_CUH + +#include + +#include "cute/layout.hpp" +#include "cute/tensor.hpp" +#include "cute/util/print.hpp" +#include "cutlass/fast_math.h" + +namespace flashinfer { + +using namespace cute; + +template +struct BlockSparseIndexedGather { + CUTE_HOST_DEVICE constexpr BlockSparseIndexedGather(IdType const* indices) : indices_(indices) {} + + template + CUTE_HOST_DEVICE constexpr IdType operator()(I i) const { + // NOTE(Zihao): there is a risk of out-of-bound access, adding boundary check here + // would degrade performance significantly. It is the user's responsibility to ensure + // that (indptr[-2] + TILE_KV) is less than the size of the indices tensor. + return indices_[i]; + } + + CUTE_HOST_DEVICE friend void print(BlockSparseIndexedGather const& s) { + cute::print("BlockSparseIndexedGather"); + } + + IdType const* indices_; +}; + +/// Custom stride object that applies a function followed by a stride +template +struct CustomStride { + CUTE_HOST_DEVICE constexpr CustomStride(Func const& func, int stride_n) + : func_(func), stride_n_(stride_n) {} + + template + CUTE_HOST_DEVICE friend auto operator*(I i, CustomStride const& s) { + // uint64_t ret; + // #if defined(__CUDA_ARCH__) + // asm("{\n\t" + // "mul.wide.u32 %0, %1, %2;\n\t" + // "}" : "=l"(ret) : "r"(s.func_(i)), "r"(s.stride_n_)); + // #else + // ret = uint64_t(s.func_(i)) * uint64_t(s.stride_n_); + // #endif + // return ret; + + // NOTE(Zihao): if the tensor is larger than 64GB ((2 ** 32) * 16byte), we use + // 64-bit multiplication to avoid overflow. Otherwise, 32-bit multiplication is + // sufficient. + // There is a 20+ TFLOPs/s gap between 32-bit and 64-bit multiplication on H100. + return uint32_t(s.func_(i)) * s.stride_n_; + } + + template + CUTE_HOST_DEVICE friend auto operator*(CustomStride const& s, I i) { + // uint64_t ret; + // #if defined(__CUDA_ARCH__) + // asm("{\n\t" + // "mul.wide.u32 %0, %1, %2;\n\t" + // "}" : "=l"(ret) : "r"(s.func_(i)), "r"(s.stride_n_)); + // #else + // ret = uint64_t(s.func_(i)) * uint64_t(s.stride_n_); + // #endif + // return ret; + + // NOTE(Zihao): if the tensor is larger than 64GB = (2 ** 32) * 16byte (16byte is the + // element size after upcasting), we use 64-bit multiplication to avoid overflow. Otherwise, + // 32-bit multiplication is sufficient. + // There is a 20+ TFLOPs/s gap between 32-bit and 64-bit multiplication on H100. + return uint32_t(s.func_(i)) * s.stride_n_; + } + + CUTE_HOST_DEVICE friend void print(CustomStride const& s) { + cute::print("BlockSparseStride{"); + print(s.func_); + cute::print(","); + print(s.stride_n_); + cute::print("}"); + } + + template + CUTE_HOST_DEVICE constexpr friend auto safe_div(CustomStride const& s, Div const& div) { + return CustomStride(s.func_, safe_div(s.stride_n_, div)); + } + + // Circumvent the requirement on make_layout that shape and stride are integral + template + CUTE_HOST_DEVICE constexpr friend auto make_layout(Shape const& shape, + CustomStride const& stride) { + return Layout(shape, stride); + } + + Func func_; + uint32_t stride_n_; +}; + +template +CUTLASS_HOST_DEVICE auto make_custom_stride_layout(int stride_n, Func&& func) { + return make_layout(make_shape(_1{}, _1{}), + make_stride(CustomStride(static_cast(func), stride_n), _1{})); +} + +/// Helper function to optionally create a block sparse gather tensor +template +CUTLASS_HOST_DEVICE auto make_block_sparse_tensor(Iterator iter, Shape const& shape, int stride_n, + Func&& func) { + Layout matrix_layout = make_identity_layout(shape); + auto offset = as_arithmetic_tuple(repeat_like(shape, _0{})); + Layout gather_layout = make_custom_stride_layout(stride_n, static_cast(func)); + + return make_tensor(iter, ComposedLayout{gather_layout, offset, matrix_layout}); +} + +} // namespace flashinfer + +namespace cute { + +template +CUTE_HOST_DEVICE constexpr auto upcast(Shape const& shape, Stride const& stride) { + if constexpr (is_tuple::value) { + return transform_layout(shape, stride, + [](auto const& s, auto const& d) { return upcast(s, d); }); + } else if constexpr (is_scaled_basis::value) { + if constexpr (Stride::mode() == I) { + return make_layout(ceil_div(shape, Int{}), ceil_div(stride, Int{})); + } else { + return make_layout(shape, stride); + } + } else { + return upcast(shape, stride); + } + + CUTE_GCC_UNREACHABLE; +} + +template +CUTE_HOST_DEVICE constexpr auto upcast( + ComposedLayout, Offset, Layout> const& layout) { + // Find index of the stride-1 mode - that is the only one that requires updating inner shape and + // offset + auto idx = + find_if(layout.layout_a().stride(), [](auto x) { return is_constant<1, decltype(x)>{}; }); + constexpr int I = decltype(idx)::value; + + // Upcast the outer layout (works as expected) + auto outer = upcast(layout.layout_a()); + + // Upcast the accumulated offset along stride-1 mode + auto offset = + as_arithmetic_tuple(replace(layout.offset(), upcast(get(layout.offset())))); + + // Upcast the inner layout's shape along stride-1 mode + auto inner = upcast(layout.layout_b().shape(), layout.layout_b().stride()); + + return composition(outer, offset, inner); +} + +} // namespace cute + +#endif // FLASHINFER_ATTENTION_HOPPER_BLOCK_SPARSE_GATHER_CUH diff --git a/include/flashinfer/attention/hopper/default_params.cuh b/include/flashinfer/attention/hopper/default_params.cuh new file mode 100644 index 0000000000000000000000000000000000000000..f2b9d2e33e072f9a1019a44bf44f40c96eb39983 --- /dev/null +++ b/include/flashinfer/attention/hopper/default_params.cuh @@ -0,0 +1,168 @@ +/* + * Copyright (c) 2024 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef FLASHINFER_ATTENTION_HOPPER_PARAMS_CUH +#define FLASHINFER_ATTENTION_HOPPER_PARAMS_CUH + +#include + +#include + +namespace flashinfer { + +template +struct SinglePrefillParams { + using DTypeQ = DTypeQ_; + using DTypeKV = DTypeKV_; + using DTypeO = DTypeO_; + using IdType = IdType_; + // The QKV matrices. + DTypeQ* q_ptr; + DTypeKV* k_ptr; + DTypeKV* v_ptr; + DTypeO* o_ptr; + float* lse_ptr; + + struct AdditionalParams { + float logits_soft_cap; + float sm_scale; + float* scale_q; + float* scale_k; + float* scale_v; + } additional_params; + + int64_t q_stride_n; + int64_t k_stride_n; + int64_t v_stride_n; + int64_t o_stride_n; + int64_t q_stride_h; + int64_t k_stride_h; + int64_t v_stride_h; + int64_t o_stride_h; + + int qo_len; + int kv_len; + int num_qo_heads; + int num_kv_heads; + int group_size; + int window_left; + + bool causal; +}; + +template +struct BatchPrefillRaggedParams { + using DTypeQ = DTypeQ_; + using DTypeKV = DTypeKV_; + using DTypeO = DTypeO_; + using IdType = IdType_; + // The QKV matrices. + DTypeQ* q_ptr; + DTypeKV* k_ptr; + DTypeKV* v_ptr; + DTypeO* o_ptr; + float* lse_ptr; + + IdType* qo_tile_indices; + IdType* qo_indptr; + IdType* kv_indptr; + IdType* qo_lens; + IdType* kv_lens; + IdType* head_indices; + IdType* work_indptr; + IdType* batch_indices; + + struct AdditionalParams { + float logits_soft_cap; + float sm_scale; + uint32_t* maybe_prefix_len_ptr; + uint16_t* maybe_token_pos_in_items_ptr; + uint32_t token_pos_in_items_len; + uint16_t* maybe_max_item_len_ptr; + } additional_params; + + int64_t q_stride_n; + int64_t k_stride_n; + int64_t v_stride_n; + int64_t o_stride_n; + int64_t q_stride_h; + int64_t k_stride_h; + int64_t v_stride_h; + int64_t o_stride_h; + int64_t nnz_qo; + int64_t nnz_kv; + + int num_qo_heads; + int num_kv_heads; + int group_size; + int window_left; + + bool causal; +}; + +template +struct BatchPrefillPagedParams { + using DTypeQ = DTypeQ_; + using DTypeKV = DTypeKV_; + using DTypeO = DTypeO_; + using IdType = IdType_; + // The QKV matrices. + DTypeQ* q_ptr; + DTypeKV* k_ptr; + DTypeKV* v_ptr; + DTypeO* o_ptr; + float* lse_ptr; + + IdType* qo_tile_indices; + IdType* qo_indptr; + IdType* kv_indptr; + IdType* kv_indices; + IdType* qo_lens; + IdType* kv_lens; + IdType* head_indices; + IdType* work_indptr; + IdType* batch_indices; + + struct AdditionalParams { + float logits_soft_cap; + float sm_scale; + uint32_t* maybe_prefix_len_ptr; + uint16_t* maybe_token_pos_in_items_ptr; + uint32_t token_pos_in_items_len; + uint16_t* maybe_max_item_len_ptr; + } additional_params; + + int64_t q_stride_n; + int64_t k_stride_n; + int64_t v_stride_n; + int64_t o_stride_n; + int64_t q_stride_h; + int64_t k_stride_h; + int64_t v_stride_h; + int64_t o_stride_h; + int64_t nnz_qo; + + int num_qo_heads; + int num_kv_heads; + int group_size; + int page_size; + int window_left; + + bool causal; +}; + +} // namespace flashinfer + +#endif // FLASHINFER_ATTENTION_HOPPER_PARAMS_CUH diff --git a/include/flashinfer/attention/hopper/epilogue.cuh b/include/flashinfer/attention/hopper/epilogue.cuh new file mode 100644 index 0000000000000000000000000000000000000000..81e43bd9a7481f519aa4617a20700603ebf0b07b --- /dev/null +++ b/include/flashinfer/attention/hopper/epilogue.cuh @@ -0,0 +1,259 @@ +/* + * Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri + * Dao. Licensed under the BSD 3-Clause. + * + * Modified by the FlashInfer team. + */ +#ifndef FLASHINFER_ATTENTION_HOPPER_EPILOGUE_CUH_ +#define FLASHINFER_ATTENTION_HOPPER_EPILOGUE_CUH_ + +#include + +#include "../../math.cuh" +#include "cute/tensor.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "named_barrier.cuh" +#include "utils.cuh" + +namespace flashinfer { + +using namespace cute; + +template +__forceinline__ __device__ void write_tiled(DTypeO* O, const TiledCopyO& tiled_copy_O, + const LayoutO& layout_O, const TileShapeO& tile_shape_O, + const SMemO& sO, int thread_idx, int qo_tile_idx, + int qo_head_idx, int qo_indptr, int64_t qo_len) { + Tensor mO = make_tensor(make_gmem_ptr(O + qo_indptr * stride<0>(layout_O)), layout_O); + Tensor gO = + get_local_tile_tensor(mO, tile_shape_O, qo_head_idx, 0, qo_len)(_, _, qo_tile_idx); // (O, D) + Tensor cO = cute::make_identity_tensor(gO.shape()); // (O, D) -> (o_idx, d_idx) + + ThrCopy thr_copy_O = tiled_copy_O.get_slice(thread_idx); + Tensor tOgO = thr_copy_O.partition_D(gO); // (CPY, CPY_O, CPY_D) + Tensor tOsO = thr_copy_O.partition_S(sO); // (CPY, CPY_O, CPY_D) + Tensor tOcO = thr_copy_O.partition_D(cO); // (CPY, CPY_O, CPY_D) + Tensor tOsOGroup = flatten_1(tOsO); // (CPY, (CPY_O, CPY_D)) + Tensor tOgOGroup = flatten_1(tOgO); // (CPY, (CPY_O, CPY_D)) + Tensor tOcOGroup = flatten_1(tOcO); // (CPY, (CPY_O, CPY_D)) + + const int qo_tile_size = get<0>(tile_shape_O); + int valid_qo_tile_size = std::min(qo_len - qo_tile_idx * qo_tile_size, qo_tile_size); + if (valid_qo_tile_size == qo_tile_size) { + copy(tiled_copy_O, tOsOGroup, tOgOGroup); + } else { + // copy if not out of bound + auto predicate_fn = [&](auto coords) { + auto s_coords = tOcOGroup(_0{}, coords); + return elem_less(get<0>(s_coords), valid_qo_tile_size); + }; + copy_if(tiled_copy_O, predicate_fn, tOsOGroup, tOgOGroup); + } +} + +template +__forceinline__ __device__ void write_O(ElemO* O, const TiledCopyO& tiled_copy_O, + const LayoutO& layout_O, const TileShapeO& tile_shape_O, + const SMemO& sO, int thread_idx, int qo_tile_idx, + int qo_head_idx, int qo_indptr, int qo_len, + int write_warp_idx) { + write_tiled(O, tiled_copy_O, layout_O, tile_shape_O, sO, thread_idx, + qo_tile_idx, qo_head_idx, qo_indptr, qo_len); +} + +template +struct CollectiveEpilogue { + using DTypeO = typename Ktraits::DTypeO; + static constexpr int CTA_Q = Ktraits::CTA_Q; + static constexpr int CTA_KV = Ktraits::CTA_KV; + static constexpr int HEAD_DIM_VO = Ktraits::HEAD_DIM_VO; + using TileShape_PDV = Shape, Int, Int>; + + static constexpr int NUM_WARPS = Ktraits::NUM_WARPS; + static constexpr int NUM_THREADS = NUM_WARPS * cutlass::NumThreadsPerWarp; + + static constexpr int NUM_COPY_THREADS = cutlass::NumThreadsPerWarpGroup; + static constexpr int NUM_MMA_THREADS = NUM_THREADS - NUM_COPY_THREADS; + + using SmemLayoutAtomO = decltype(cutlass::gemm::collective::detail::ss_smem_selector< + GMMA::Major::K, DTypeO, decltype(cute::get<0>(TileShape_PDV{})), + decltype(cute::get<1>(TileShape_PDV{}))>()); + using SmemLayoutO = decltype(tile_to_shape(SmemLayoutAtomO{}, select<0, 1>(TileShape_PDV{}))); + + using SmemCopyAtomO = Copy_Atom; + using SharedStorage = cute::array_aligned>; + + using ShapeT = cute::Shape; + using StrideT = cute::Shape; + using LayoutT = cute::Layout; + + using ShapeLseT = cute::Shape; + using StrideLseT = cute::Shape<_1, int64_t>; + using LayoutLseT = cute::Layout; + + using GmemTiledCopyOTMA = cute::SM90_TMA_STORE; + using TMA_O = decltype(make_tma_copy( + GmemTiledCopyOTMA{}, + make_tensor(make_gmem_ptr(static_cast(nullptr)), ShapeT{}, StrideT{}), SmemLayoutO{}, + select<0, 1>(TileShape_PDV{}), _1{})); // no mcast for O + + static constexpr int VEC_SIZE = cute::ceil_div(128, sizeof_bits_v); + static_assert(HEAD_DIM_VO % VEC_SIZE == 0); + static constexpr int NUM_THREADS_PER_ROW = HEAD_DIM_VO / VEC_SIZE; + static_assert(NUM_MMA_THREADS % NUM_THREADS_PER_ROW == 0); + static constexpr int NUM_ROWS = NUM_MMA_THREADS / NUM_THREADS_PER_ROW; + using TiledCopyOAtom = cute::Copy_Atom, DTypeO>; + using TiledCopyOThrLayout = decltype(cute::make_layout( + cute::make_shape(Int{}, Int{}), LayoutRight{})); + using TiledCopyOValLayout = + decltype(cute::make_layout(cute::make_shape(_1{}, Int{}), LayoutRight{})); + using TiledCopyO = + decltype(make_tiled_copy(TiledCopyOAtom{}, TiledCopyOThrLayout{}, // Thr layout + TiledCopyOValLayout{} // Val layout + )); + + // used for rmem -> smem O copy in fp8 kernel to undo column permutation + using ThreadLayoutrO = Layout, _4, _1>, Stride<_4, _32, _1, _0>>; + using ValueLayoutrO = Layout, Int>, + Stride<_0, _2, Stride<_4, _1>, _8>>; + using TiledCopyrO = decltype(make_tiled_copy(Copy_Atom, DTypeO>{}, + ThreadLayoutrO{}, ValueLayoutrO{})); + using TiledCopyShaperO = Shape<_8, Int, _16, Int>; + using SmemLayoutrO = decltype(composition(SmemLayoutO{}, Layout{})); + + // Host side kernel arguments + struct Arguments { + DTypeO* O_ptr; + LayoutT const layout_O; + float* lse_ptr; + LayoutLseT const layout_LSE; + }; + + // Device side kernel params + struct Params { + DTypeO* O_ptr; + LayoutT const layout_O; + float* lse_ptr; + LayoutLseT const layout_LSE; + }; + + static Params to_underlying_arguments(Arguments const& args) { + Tensor mO = make_tensor(make_gmem_ptr(args.O_ptr), args.layout_O); + return {args.O_ptr, args.layout_O, args.lse_ptr, args.layout_LSE}; + } + + /// Issue Tma Descriptor Prefetch -- ideally from a single thread for best performance + CUTLASS_DEVICE + static void prefetch_tma_descriptors(Params const& epilogue_params) {} + + template + CUTLASS_DEVICE void store(Params const& epilogue_params, FrgTensorO const& tOrO, + FrgTensorLSE const& lse, SharedStorage& shared_storage, + TiledMma tiled_mma, int thread_idx, BlockCoord const& block_coord) { + auto [qo_tile_idx, qo_head_idx, kv_head_idx, qo_indptr, kv_indptr, qo_len, kv_len, batch_idx] = + block_coord; + Tensor sO = make_tensor(make_smem_ptr(shared_storage.smem_o.data()), SmemLayoutO{}); + auto smem_tiled_copy_O = make_tiled_copy_C(SmemCopyAtomO{}, tiled_mma); + auto smem_thr_copy_O = smem_tiled_copy_O.get_thread_slice(thread_idx); + + Tensor tOrO_out = convert_type(tOrO); + Tensor tOrO_retile = smem_thr_copy_O.retile_S(tOrO_out); // ((Atom,AtomNum), MMA_M, MMA_N) + Tensor tOsO = smem_thr_copy_O.partition_D(sO); // ((Atom,AtomNum),PIPE_M,PIPE_N) + + // Make sure all WGs have finished reading V + cutlass::arch::NamedBarrier::sync(NUM_MMA_THREADS, + /*id=*/static_cast(NamedBarriers::kValueEmpty)); + cute::copy(smem_tiled_copy_O, tOrO_retile, tOsO); + cutlass::arch::fence_view_async_shared(); // ensure smem writes are visible to TMA + cutlass::arch::NamedBarrier::arrive(NUM_MMA_THREADS + Ktraits::NUM_PRODUCER_THREADS, + cutlass::arch::ReservedNamedBarriers::EpilogueBarrier); + + Tensor mLSE = make_tensor(make_gmem_ptr(epilogue_params.lse_ptr), epilogue_params.layout_LSE); + Tensor gLSE = get_lse_local_tile_tensor(mLSE, Shape>{}, qo_head_idx, qo_indptr, + qo_len)(_, qo_tile_idx); + Tensor cO = cute::make_identity_tensor(select<0, 1>(TileShape_PDV{})); + auto thread_mma = tiled_mma.get_thread_slice(thread_idx); + Tensor tOcO = thread_mma.partition_C(cO); // (MMA,MMA_M,MMA_K) + static_assert(decltype(size<0, 0>(tOcO))::value == 2); + static_assert(decltype(size<0, 1>(tOcO))::value == 2); + // tOcO has shape ((2, 2, V), MMA_M, MMA_K), we only take only the row indices. + Tensor tOcO_row = tOcO(make_coord(_0{}, _, _0{}), _, _0{}); + CUTE_STATIC_ASSERT_V(size(lse) == size(tOcO_row)); // MMA_M + if (epilogue_params.lse_ptr) { // don't write to LSE if it's nullptr + if (get<1>(tOcO_row(_0{})) == 0) { +#pragma unroll + for (int mi = 0; mi < size(lse); ++mi) { + const int row = get<0>(tOcO_row(mi)); + if (row < qo_len - qo_tile_idx * CTA_Q) { + gLSE(row) = lse(mi); + } + } + } + } + + int write_warp_idx = NUM_WARPS - 1; + if (cutlass::canonical_warp_idx_sync() == write_warp_idx) { + cutlass::arch::NamedBarrier::sync(NUM_MMA_THREADS + Ktraits::NUM_PRODUCER_THREADS, + cutlass::arch::ReservedNamedBarriers::EpilogueBarrier); + } + TiledCopyO gmem_tiled_copy_O; + write_O(epilogue_params.O_ptr, gmem_tiled_copy_O, epilogue_params.layout_O, + select<0, 1>(TileShape_PDV{}), sO, thread_idx, qo_tile_idx, + qo_head_idx, qo_indptr, qo_len, write_warp_idx); + } + + CUTLASS_DEVICE void store_tail() { + // tma_store_wait<0>(); + } + + // Write 0 to output and -inf to LSE + template + CUTLASS_DEVICE void store_zero(Params const& epilogue_params, SharedStorage& shared_storage, + int thread_idx, BlockCoord const& block_coord) { + auto [qo_tile_idx, qo_head_idx, kv_head_idx, qo_indptr, kv_indptr, qo_len, kv_len, batch_idx] = + block_coord; + Tensor mO = make_tensor(make_gmem_ptr(epilogue_params.O_ptr), epilogue_params.layout_O); + Tensor gO = get_local_tile_tensor(mO, select<0, 1>(TileShape_PDV{}), qo_head_idx, qo_indptr, + qo_len)(_, _, qo_tile_idx); // (O, D) + Tensor cO = cute::make_identity_tensor(gO.shape()); // (O, D) -> (o_idx, d_idx) + Tensor mLSE = make_tensor(make_gmem_ptr(epilogue_params.lse_ptr), epilogue_params.layout_LSE); + Tensor gLSE = get_lse_local_tile_tensor(mLSE, Shape>{}, qo_head_idx, qo_indptr, + qo_len)(_, qo_tile_idx); + + TiledCopyO tiled_copy_O; + auto thr_copy_O = tiled_copy_O.get_thread_slice(thread_idx); + Tensor tOgO = thr_copy_O.partition_D(gO); // (CPY, CPY_O, CPY_D) + Tensor tOrO = make_fragment_like(tOgO); // (CPY, CPY_O, CPY_D) + clear(tOrO); + Tensor tOcO = thr_copy_O.partition_D(cO); // (CPY, CPY_O, CPY_D) + Tensor tOgOGroup = flatten_1(tOgO); // (CPY, (CPY_O, CPY_D)) + Tensor tOrOGroup = flatten_1(tOrO); // (CPY, (CPY_O, CPY_D)) + Tensor tOcOGroup = flatten_1(tOcO); // (CPY, (CPY_O, CPY_D)) + + const int qo_tile_size = get<0>(TileShape_PDV{}); + int valid_qo_tile_size = std::min(qo_len - qo_tile_idx * qo_tile_size, qo_tile_size); + if (valid_qo_tile_size == qo_tile_size) { + copy(tiled_copy_O, tOrOGroup, tOgOGroup); + } else { + auto predicate_fn = [&](auto coords) { + auto s_coords = tOcOGroup(_0{}, coords); + return elem_less(get<0>(s_coords), valid_qo_tile_size); + }; + copy_if(tiled_copy_O, predicate_fn, tOrOGroup, tOgOGroup); + } + + static_assert(CTA_Q <= NUM_MMA_THREADS); + if (epilogue_params.lse_ptr) { // don't write to LSE if it's nullptr + if (thread_idx < qo_len - qo_tile_idx * CTA_Q) { + gLSE(thread_idx) = -math::inf; + } + } + } +}; + +} // namespace flashinfer + +#endif // FLASHINFER_ATTENTION_HOPPER_EPILOGUE_CUH_ diff --git a/include/flashinfer/attention/hopper/kernel_traits.cuh b/include/flashinfer/attention/hopper/kernel_traits.cuh new file mode 100644 index 0000000000000000000000000000000000000000..2ac599ca64076c86ea2a7f87768a301ea0dce876 --- /dev/null +++ b/include/flashinfer/attention/hopper/kernel_traits.cuh @@ -0,0 +1,123 @@ +/* + * Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri + * Dao. Licensed under the BSD 3-Clause. + * + * Modified by the FlashInfer team. + */ +#ifndef FLASHINFER_ATTENTION_HOPPER_KERNEL_TRAITS_CUH_ +#define FLASHINFER_ATTENTION_HOPPER_KERNEL_TRAITS_CUH_ + +#include + +#include "../../cutlass_utils.cuh" +#include "cute/algorithm/copy.hpp" +#include "cute/atom/mma_atom.hpp" +#include "cutlass/cutlass.h" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/layout/layout.h" +#include "cutlass/numeric_types.h" +#include "cutlass/pipeline/pipeline.hpp" + +namespace flashinfer { + +using namespace cute; + +template +struct SharedStorageQKVO { + cute::array_aligned> smem_q; + cute::array_aligned> smem_k; + union { + cute::array_aligned> smem_v; + cute::array_aligned> smem_o; + }; + struct { + cutlass::arch::ClusterTransactionBarrier barrier_Q; + cutlass::arch::ClusterBarrier barrier_O; + typename MainloopPipeline::SharedStorage pipeline_k; + typename MainloopPipeline::SharedStorage pipeline_v; + }; +}; + +template +struct AttentionKernelTraits { + using AttentionVariant = AttentionVariant_; + + using DTypeQ = DTypeQ_; + using DTypeKV = DTypeKV_; + using DTypeO = DTypeO_; + using IdType = IdType_; + using DTypeQKAccum = float; + + static constexpr int CTA_Q = CTA_Q_; + static_assert(CTA_Q % 64 == 0); + static constexpr int CTA_KV = CTA_KV_; + static constexpr int HEAD_DIM_QK = HEAD_DIM_QK_; + static constexpr int HEAD_DIM_VO = HEAD_DIM_VO_; + static_assert(HEAD_DIM_QK % 32 == 0); + static_assert(HEAD_DIM_VO % 32 == 0); + + static constexpr int NUM_WARPS = ((CTA_Q / 64) + 1) * 4; + static constexpr int NUM_THREADS = NUM_WARPS * cutlass::NumThreadsPerWarp; + // NOTE(Zihao): the following constant should only be used when TMA is enabled, + // where only one warp inside a warp group is used for TMA. + static constexpr int NUM_PRODUCER_THREADS = cutlass::NumThreadsPerWarp; + + using TileShape_QKD = Shape, Int, Int>; + using TileShape_PDV = Shape, Int, Int>; + + static constexpr int NUM_STAGES = NUM_STAGES_; + + using AtomLayoutQKD = Layout, _1, _1>>; + using TiledMmaQK = decltype(cute::make_tiled_mma( + cute::GMMA::ss_op_selector(), AtomLayoutQKD{})); + using TiledMmaPV = decltype(cute::make_tiled_mma( + cute::GMMA::rs_op_selector(), + AtomLayoutQKD{})); + + static constexpr int NUM_MMA_THREADS = size(TiledMmaQK{}); + + using SmemLayoutAtomQ = decltype(cutlass::gemm::collective::detail::ss_smem_selector< + GMMA::Major::K, DTypeQ, decltype(cute::get<0>(TileShape_QKD{})), + decltype(cute::get<2>(TileShape_QKD{}))>()); + using SmemLayoutQ = decltype(tile_to_shape(SmemLayoutAtomQ{}, select<0, 2>(TileShape_QKD{}))); + + using SmemLayoutAtomK = decltype(cutlass::gemm::collective::detail::ss_smem_selector< + GMMA::Major::K, DTypeKV, decltype(cute::get<1>(TileShape_QKD{})), + decltype(cute::get<2>(TileShape_QKD{}))>()); + using SmemLayoutK = decltype(tile_to_shape( + SmemLayoutAtomK{}, + make_shape(shape<1>(TileShape_QKD{}), shape<2>(TileShape_QKD{}), Int{}))); + + using SmemLayoutAtomV = decltype(cutlass::gemm::collective::detail::ss_smem_selector< + GMMA::Major::K, DTypeKV, decltype(cute::get<2>(TileShape_PDV{})), + decltype(cute::get<1>(TileShape_PDV{}))>()); + using SmemLayoutV = decltype(tile_to_shape( + SmemLayoutAtomV{}, + make_shape(get<2>(TileShape_PDV{}), get<1>(TileShape_PDV{}), Int{}))); + + // Note this is the transpose in terms of the view, not in terms of memory. + using SmemLayoutVt = decltype(composition( + SmemLayoutV{}, make_ordered_layout(make_shape(get<1>(TileShape_PDV{}), + get<2>(TileShape_PDV{}), Int{}), + Step<_2, _1, _3>{}))); + + using SmemLayoutAtomO = decltype(cutlass::gemm::collective::detail::ss_smem_selector< + GMMA::Major::K, DTypeO, decltype(cute::get<0>(TileShape_PDV{})), + decltype(cute::get<1>(TileShape_PDV{}))>()); + using SmemLayoutO = decltype(tile_to_shape(SmemLayoutAtomO{}, select<0, 1>(TileShape_PDV{}))); + using MainloopPipeline = + std::conditional_t, + typename cutlass::PipelineAsync>; + using PipelineState = typename cutlass::PipelineState; + + using SharedStorage = SharedStorageQKVO; +}; + +} // namespace flashinfer + +#endif // FLASHINFER_ATTENTION_HOPPER_KERNEL_TRAITS_CUH_ diff --git a/include/flashinfer/attention/hopper/mainloop.cuh b/include/flashinfer/attention/hopper/mainloop.cuh new file mode 100644 index 0000000000000000000000000000000000000000..e5bf4ffb9f3aab8984c7daa3c453dd6653890830 --- /dev/null +++ b/include/flashinfer/attention/hopper/mainloop.cuh @@ -0,0 +1,269 @@ +/* + * Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri + * Dao. Licensed under the BSD 3-Clause. + * + * Modified by the FlashInfer team. + */ +#ifndef FLASHINFER_ATTENTION_HOPPER_MAINLOOP_CUH_ +#define FLASHINFER_ATTENTION_HOPPER_MAINLOOP_CUH_ + +#include +#include +#include +#include + +#include "../../math.cuh" +#include "cute/tensor.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/pipeline/pipeline.hpp" +#include "mainloop_mma.cuh" +#include "named_barrier.cuh" +#include "utils.cuh" + +namespace flashinfer { + +using namespace cute; + +template +struct CollectiveMainloop { + using DTypeQ = typename Ktraits::DTypeQ; + using DTypeKV = typename Ktraits::DTypeKV; + using TileShape_QKD = typename Ktraits::TileShape_QKD; + using TileShape_PDV = typename Ktraits::TileShape_PDV; + static constexpr int CTA_Q = get<0>(TileShape_QKD{}); + static constexpr int CTA_KV = get<1>(TileShape_QKD{}); + + static constexpr int NUM_STAGES = Ktraits::NUM_STAGES; + static constexpr int NUM_MMA_THREADS = Ktraits::NUM_MMA_THREADS; + static constexpr int HEAD_DIM_QK = Ktraits::HEAD_DIM_QK; + static constexpr int HEAD_DIM_VO = Ktraits::HEAD_DIM_VO; + + using GmemTiledCopyQ = cute::SM90_TMA_LOAD; + using GmemTiledCopyKV = cute::SM90_TMA_LOAD; + + using SmemLayoutQ = typename Ktraits::SmemLayoutQ; + using SmemLayoutK = typename Ktraits::SmemLayoutK; + using SmemLayoutV = typename Ktraits::SmemLayoutV; + using SmemLayoutVt = typename Ktraits::SmemLayoutVt; + + using ShapeT = cute::Shape; + using StrideT = cute::Shape; // (N, D, H) + using LayoutT = cute::Layout; + + using ShapeLseT = cute::Shape; + using StrideLseT = cute::Shape<_1, int64_t>; + using LayoutLseT = cute::Layout; + + using TMA_Q = decltype(make_tma_copy( + GmemTiledCopyQ{}, + make_tensor(make_gmem_ptr(static_cast(nullptr)), + repeat_like(StrideT{}, int32_t(0)), StrideT{}), + SmemLayoutQ{}, select<0, 2>(TileShape_QKD{}), _1{})); // no mcast for Q + + using TMA_K = decltype(make_tma_copy( + GmemTiledCopyKV{}, + make_tensor(make_gmem_ptr(static_cast(nullptr)), + repeat_like(StrideT{}, int32_t(0)), StrideT{}), + take<0, 2>(SmemLayoutK{}), select<1, 2>(TileShape_QKD{}), _1{})); // no mcast + + using TMA_V = decltype(make_tma_copy( + GmemTiledCopyKV{}, + make_tensor(make_gmem_ptr(static_cast(nullptr)), + repeat_like(StrideT{}, int32_t(0)), StrideT{}), + take<0, 2>(SmemLayoutV{}), select<2, 1>(TileShape_PDV{}), _1{})); // no mcast + + static constexpr bool USE_TMA_LOAD_KV = true; + using MainloopPipeline = typename Ktraits::MainloopPipeline; + using PipelineParams = typename MainloopPipeline::Params; + using PipelineState = typename MainloopPipeline::PipelineState; + + // Set the bytes transferred in this TMA transaction (may involve multiple issues) + static constexpr uint32_t TmaTransactionBytesQ = + static_cast(size(SmemLayoutQ{}) * cutlass::sizeof_bits_v / 8); + static constexpr uint32_t TmaTransactionBytesK = + static_cast(size(take<0, 2>(SmemLayoutK{})) * cutlass::sizeof_bits_v / 8); + static constexpr uint32_t TmaTransactionBytesV = + static_cast(size(take<0, 2>(SmemLayoutV{})) * cutlass::sizeof_bits_v / 8); + + // Whether use scheduler barrier or hardware warp scheduler, using heuristic based on data type + // and head dim + static constexpr bool UseSchedulerBarrier = + cutlass::sizeof_bits_v == 8 ? HEAD_DIM_VO >= 128 : HEAD_DIM_VO <= 128; + using WarpScheduler = WarpScheduler; + + // Host side kernel arguments + struct Arguments { + DTypeQ const* Q_ptr; + LayoutT layout_Q; + DTypeKV const* K_ptr; + LayoutT layout_K; + DTypeKV const* V_ptr; + LayoutT layout_V; + int window_left; + AdditionalParams additional_params; + }; + + // Device side kernel params + struct Params { + LayoutT layout_Q; + LayoutT layout_K; + LayoutT layout_V; + TMA_Q tma_load_Q; + TMA_K tma_load_K; + TMA_V tma_load_V; + int window_left; + AdditionalParams additional_params; + }; + + static Params to_underlying_arguments(Arguments const& args) { + Tensor mQ = make_tensor(make_gmem_ptr(args.Q_ptr), args.layout_Q); + TMA_Q tma_load_Q = make_tma_copy(GmemTiledCopyQ{}, mQ, SmemLayoutQ{}, + select<0, 2>(TileShape_QKD{}), _1{}); // no mcast for Q + Tensor mK = make_tensor(make_gmem_ptr(args.K_ptr), args.layout_K); + TMA_K tma_load_K = make_tma_copy(GmemTiledCopyKV{}, mK, SmemLayoutK{}(_, _, _0{}), + select<1, 2>(TileShape_QKD{}), _1{}); // no mcast + Tensor mV = make_tensor(make_gmem_ptr(args.V_ptr), args.layout_V); + TMA_V tma_load_V = make_tma_copy(GmemTiledCopyKV{}, mV, SmemLayoutV{}(_, _, _0{}), + select<2, 1>(TileShape_PDV{}), _1{}); // no mcast + return {args.layout_Q, args.layout_K, args.layout_V, tma_load_Q, + tma_load_K, tma_load_V, args.window_left, args.additional_params}; + } + + /// Issue Tma Descriptor Prefetch -- ideally from a single thread for best performance + CUTLASS_DEVICE + static void prefetch_tma_descriptors(Params const& mainloop_params) { + cute::prefetch_tma_descriptor(mainloop_params.tma_load_Q.get_tma_descriptor()); + cute::prefetch_tma_descriptor(mainloop_params.tma_load_K.get_tma_descriptor()); + cute::prefetch_tma_descriptor(mainloop_params.tma_load_V.get_tma_descriptor()); + } + + CUTLASS_DEVICE + int get_num_kv_tiles(Params const& mainloop_params, int q_tile_idx, const int qo_len, + const int kv_len) { + static constexpr int CTA_Q = get<0>(TileShape_QKD{}); + static constexpr int CTA_KV = get<1>(TileShape_QKD{}); + int num_kv_tiles = cute::ceil_div(kv_len, CTA_KV); + if constexpr (CAUSAL) { + num_kv_tiles = std::min(num_kv_tiles, + cute::ceil_div((q_tile_idx + 1) * CTA_Q + kv_len - qo_len, CTA_KV)); + } + + return num_kv_tiles; + } + + template + CUTLASS_DEVICE void load(Params const& mainloop_params, MainloopPipeline pipeline_k, + MainloopPipeline pipeline_v, PipelineState& smem_pipe_write_k, + PipelineState& smem_pipe_write_v, SharedStorage& shared_storage, + Scheduler& scheduler, typename Scheduler::Params const& scheduler_params, + typename Scheduler::WorkTileInfo& work_tile_info, + BlockCoord const& block_coord, int work_idx) { + Tensor sQ = make_tensor(make_smem_ptr(shared_storage.smem_q.data()), SmemLayoutQ{}); + Tensor sK = make_tensor(make_smem_ptr(shared_storage.smem_k.data()), SmemLayoutK{}); + Tensor sV = make_tensor(make_smem_ptr(shared_storage.smem_v.data()), SmemLayoutV{}); + + Tensor mQ = mainloop_params.tma_load_Q.get_tma_tensor(mainloop_params.layout_Q.shape()); + Tensor mK = mainloop_params.tma_load_K.get_tma_tensor(mainloop_params.layout_K.shape()); + Tensor mV = mainloop_params.tma_load_V.get_tma_tensor(mainloop_params.layout_V.shape()); + + auto [q_tile_idx, qo_head_idx, kv_head_idx, qo_indptr, kv_indptr, qo_len, kv_len, batch_idx] = + block_coord; + + // Prepare the TMA loads + Tensor gQ = get_local_tile_tensor(mQ, select<0, 2>(TileShape_QKD{}), qo_head_idx, qo_indptr, + qo_len)(_, _, q_tile_idx); // (Q, D) + Tensor gK = get_local_tile_tensor(mK, select<1, 2>(TileShape_QKD{}), kv_head_idx, kv_indptr, + kv_len); // (K, D, _) + Tensor gV = get_local_tile_tensor(mV, select<2, 1>(TileShape_PDV{}), kv_head_idx, kv_indptr, + kv_len); // (K, D, _) + + Tensor sQ_x = make_tensor(sQ.data(), make_layout(sQ.layout(), Layout<_1>{})); + Tensor gQ_x = make_tensor(gQ.data(), make_layout(gQ.layout(), Layout<_1>{})); + auto [tQgQ, tQsQ] = + tma_partition(mainloop_params.tma_load_Q, _0{}, Layout<_1>{}, group_modes<0, 2>(sQ_x), + group_modes<0, 2>(gQ_x)); // (TMA), (TMA) + auto [tKgK, tKsK] = + tma_partition(mainloop_params.tma_load_K, _0{}, Layout<_1>{}, group_modes<0, 2>(sK), + group_modes<0, 2>(gK)); // (TMA, k), (TMA, PIPE) + auto [tVgV, tVsV] = + tma_partition(mainloop_params.tma_load_V, _0{}, Layout<_1>{}, group_modes<0, 2>(sV), + group_modes<0, 2>(gV)); // (TMA, k), (TMA, PIPE) + + int num_kv_tiles = get_num_kv_tiles(mainloop_params, q_tile_idx, qo_len, kv_len); + int kv_tile_idx = num_kv_tiles - 1; + int swa_begin_kv_tile_idx = 0; + if constexpr (LEFT_SLIDING_WINDOW) { + swa_begin_kv_tile_idx = get_swa_begin_kv_tile_idx(mainloop_params.window_left, + q_tile_idx, qo_len, kv_len); + } + + int lane_predicate = cute::elect_one_sync(); + if (lane_predicate) { + pipeline_k.producer_acquire(smem_pipe_write_k); + copy(mainloop_params.tma_load_K.with(*pipeline_k.producer_get_barrier(smem_pipe_write_k), + /*mcast_mask=*/0), + tKgK(_, kv_tile_idx), tKsK(_, smem_pipe_write_k.index())); + ++smem_pipe_write_k; + } + + // Wait for the MMA warpgroups to say that smem_q is ready + cutlass::arch::NamedBarrier::sync(NUM_MMA_THREADS + Ktraits::NUM_PRODUCER_THREADS, + static_cast(NamedBarriers::kQueryEmpty)); + + if (lane_predicate) { + shared_storage.barrier_Q.arrive_and_expect_tx(TmaTransactionBytesQ); + copy(mainloop_params.tma_load_Q.with( + reinterpret_cast( + shared_storage.barrier_Q), + /*mcast_mask=*/0), + tQgQ, tQsQ); + } + + // Wait for warp 1 to signal that smem_v are ready and V can be copied from gmem + // Need ClusterBarrier, not just NamedBarrier. Otherwise we might have CTA 0 finishing the + // TMA store on O first, call TMA multicast load on V, before CTA 1 can finishing TMA store on + // O. + shared_storage.barrier_O.wait((work_idx + 1) % 2); + + if (lane_predicate) { +#pragma unroll 2 + for (; kv_tile_idx > swa_begin_kv_tile_idx; --kv_tile_idx) { + pipeline_k.producer_acquire(smem_pipe_write_k); + copy(mainloop_params.tma_load_K.with(*pipeline_k.producer_get_barrier(smem_pipe_write_k), + /*mcast_mask=*/0), + tKgK(_, kv_tile_idx - 1), tKsK(_, smem_pipe_write_k.index())); + ++smem_pipe_write_k; + pipeline_v.producer_acquire(smem_pipe_write_v); + copy(mainloop_params.tma_load_V.with(*pipeline_v.producer_get_barrier(smem_pipe_write_v), + /*mcast_mask=*/0), + tVgV(_, kv_tile_idx), tVsV(_, smem_pipe_write_v.index())); + ++smem_pipe_write_v; + } + } + scheduler.prefetch_next_work(scheduler_params, work_tile_info); + if (lane_predicate) { + pipeline_v.producer_acquire(smem_pipe_write_v); + copy(mainloop_params.tma_load_V.with(*pipeline_v.producer_get_barrier(smem_pipe_write_v), + /*mcast_mask=*/0), + tVgV(_, kv_tile_idx), tVsV(_, smem_pipe_write_v.index())); + ++smem_pipe_write_v; + } + scheduler.broadcast_next_work(work_tile_info); + } + + CUTLASS_DEVICE void load_tail(MainloopPipeline pipeline_k, MainloopPipeline pipeline_v, + PipelineState& smem_pipe_write_k, + PipelineState& smem_pipe_write_v) { + int lane_predicate = cute::elect_one_sync(); + int warp_idx_in_warpgroup = __shfl_sync(0xffffffff, (threadIdx.x / 32) % 4, 0); + if (warp_idx_in_warpgroup == 0 && lane_predicate) { + pipeline_k.producer_tail(smem_pipe_write_k); + pipeline_v.producer_tail(smem_pipe_write_v); + } + } +}; + +} // namespace flashinfer + +#endif // FLASHINFER_ATTENTION_HOPPER_MAINLOOP_CUH_ diff --git a/include/flashinfer/attention/hopper/mainloop_mma.cuh b/include/flashinfer/attention/hopper/mainloop_mma.cuh new file mode 100644 index 0000000000000000000000000000000000000000..0348e7b8dcaef98684fc39fe45475b2761881e94 --- /dev/null +++ b/include/flashinfer/attention/hopper/mainloop_mma.cuh @@ -0,0 +1,330 @@ +/* + * Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri + * Dao. Licensed under the BSD 3-Clause. + * + * Modified by the FlashInfer team. + */ +#ifndef FLASHINFER_ATTENTION_HOPPER_MAINLOOP_MMA_CUH_ +#define FLASHINFER_ATTENTION_HOPPER_MAINLOOP_MMA_CUH_ + +#include +#include +#include +#include + +namespace flashinfer { + +template +CUTLASS_DEVICE void mma_f16( + const Params& mainloop_params, AttentionVariant& variant, MainloopPipeline pipeline_k, + MainloopPipeline pipeline_v, PipelineState& smem_pipe_read_k, PipelineState& smem_pipe_read_v, + FrgTensorO& tOrO, AttentionUpdater& attention_updater, int kv_tile_idx_count, + int swa_begin_kv_tile_idx, int swa_end_kv_tile_idx, int thread_idx, int work_idx, + int q_tile_idx, SharedStorage& shared_storage, const int32_t qo_len, const int32_t kv_len, + const int32_t qo_head_idx, const int32_t kv_head_idx, const uint32_t prefix_len, + uint16_t* token_pos_in_items, const int num_kv_tiles_outside_items_window = 0, + const int num_kv_tiles_prefix = 0) { + using DTypeQ = typename Ktraits::DTypeQ; + using DTypeKV = typename Ktraits::DTypeKV; + using IdType = typename Ktraits::IdType; + using TileShape_QKD = typename Ktraits::TileShape_QKD; + static constexpr int NUM_MMA_THREADS = Ktraits::NUM_MMA_THREADS; + using SmemLayoutQ = typename Ktraits::SmemLayoutQ; + using SmemLayoutK = typename Ktraits::SmemLayoutK; + using SmemLayoutV = typename Ktraits::SmemLayoutV; + using SmemLayoutVt = typename Ktraits::SmemLayoutVt; + static_assert(is_rmem::value, "O tensor must be rmem resident."); + + static constexpr int CTA_Q = get<0>(TileShape_QKD{}); + static constexpr int CTA_KV = get<1>(TileShape_QKD{}); + + Tensor sQ = make_tensor(make_smem_ptr(shared_storage.smem_q.data()), SmemLayoutQ{}); + Tensor sK = make_tensor(make_smem_ptr(shared_storage.smem_k.data()), SmemLayoutK{}); + Tensor sVt = make_tensor(make_smem_ptr(shared_storage.smem_v.data()), SmemLayoutVt{}); + + typename Ktraits::TiledMmaQK tiled_mma_qk; + typename Ktraits::TiledMmaPV tiled_mma_pv; + auto threadMmaQK = tiled_mma_qk.get_thread_slice(thread_idx); + auto threadMmaPV = tiled_mma_pv.get_thread_slice(thread_idx); + + Tensor tSrQ = threadMmaQK.partition_fragment_A(sQ); + Tensor tSrK = threadMmaQK.partition_fragment_B(sK); + Tensor tOrV = threadMmaPV.partition_fragment_B(sVt); + + auto consumer_wait = [](auto& pipeline, auto& smem_pipe_read) { + auto barrier_token = pipeline.consumer_try_wait(smem_pipe_read); + pipeline.consumer_wait(smem_pipe_read, barrier_token); + }; + + tiled_mma_pv.accumulate_ = GMMA::ScaleOut::Zero; + int kv_tile_idx = kv_tile_idx_count - 1; + + cutlass::ConsumerToken barrier_token = + static_cast(shared_storage.barrier_Q.try_wait(work_idx % 2)); + if (barrier_token == cutlass::BarrierStatus::WaitAgain) { + shared_storage.barrier_Q.wait(work_idx % 2); + } + + Tensor tSrS = partition_fragment_C(tiled_mma_qk, select<0, 1>(TileShape_QKD{})); + consumer_wait(pipeline_k, smem_pipe_read_k); + + WarpScheduler::barrier_sync(); + gemm(tiled_mma_qk, tSrQ, tSrK(_, _, _, smem_pipe_read_k.index()), + tSrS); + WarpScheduler::barrier_arrive(); + + if (work_idx != 0) { + int lane_predicate = cute::elect_one_sync(); + if (cutlass::canonical_warp_idx_sync() == Ktraits::NUM_WARPS - 1 && lane_predicate) { +#pragma unroll + for (uint32_t cta_id = 0; cta_id < 1; ++cta_id) { + shared_storage.barrier_O.arrive(cta_id, lane_predicate); + } + } + } + warpgroup_wait<0>(); + pipeline_k.consumer_release(smem_pipe_read_k); + ++smem_pipe_read_k; + + auto col_limit_right = [&](int qo_idx) { return qo_idx + 1 + kv_len - qo_len; }; + auto col_limit_left = [&](int qo_idx) { + return qo_idx + kv_len - qo_len - mainloop_params.window_left; + }; + auto mask_multi_item_scoring = [&](decltype(tSrS)& tSrS, int i, int qo_idx, int kv_idx) { + const uint32_t idx_in_original_seq = qo_idx + kv_len - qo_len; + const bool out_of_boundary = + kv_idx > idx_in_original_seq || (kv_idx >= std::min(kv_len, col_limit_right(qo_idx))); + const bool is_prefix = idx_in_original_seq < prefix_len; + uint16_t token_pos_in_items_regs = 0; + // Only access idx_in_original_seq >= prefix_len && idx_in_original_seq < kv_len to avoid + // out-of-bounds memory access + if (idx_in_original_seq >= prefix_len & idx_in_original_seq < kv_len) { + token_pos_in_items_regs = __ldca(token_pos_in_items + idx_in_original_seq - prefix_len); + } + if (out_of_boundary || is_prefix) { + tSrS(i) = out_of_boundary ? (AttentionUpdater::fill_value) : tSrS(i); + } else { + tSrS(i) = (kv_idx < prefix_len | (idx_in_original_seq < kv_idx + token_pos_in_items_regs)) + ? tSrS(i) + : (AttentionUpdater::fill_value); + } + }; + auto mask_multi_item_scoring_assume_in_bound = [&](decltype(tSrS)& tSrS, int i, int qo_idx, + int kv_idx) { + const uint32_t idx_in_original_seq = qo_idx + kv_len - qo_len; + const bool is_prefix = idx_in_original_seq < prefix_len; + if (is_prefix) { + tSrS(i) = AttentionUpdater::fill_value; + } else { + uint16_t token_pos_in_items_regs = 0; + // Only access idx_in_original_seq >= prefix_len && idx_in_original_seq < kv_len to avoid + // out-of-bounds memory access + if (idx_in_original_seq >= prefix_len & idx_in_original_seq < kv_len) { + token_pos_in_items_regs = __ldca(token_pos_in_items + idx_in_original_seq - prefix_len); + } + + tSrS(i) = (kv_idx < prefix_len | (idx_in_original_seq < kv_idx + token_pos_in_items_regs)) + ? tSrS(i) + : (AttentionUpdater::fill_value); + } + }; + auto kv_tile_idx_decrement = [&](int kv_tile_idx) { + int result = kv_tile_idx - 1; + if constexpr (MULTIITEMSCORING) { + if ((kv_tile_idx == num_kv_tiles_outside_items_window) & + (kv_tile_idx >= num_kv_tiles_prefix)) { + result = num_kv_tiles_prefix - 1; + } + } + return result; + }; + { + Tensor cS = cute::make_identity_tensor(select<0, 1>(TileShape_QKD{})); + Tensor tScS = threadMmaQK.partition_C(cS); +#pragma unroll + for (int i = 0; i < size(tSrS); ++i) { + int qo_idx = get<0>(tScS(i)) + q_tile_idx * CTA_Q; + int kv_idx = get<1>(tScS(i)) + kv_tile_idx * CTA_KV; + tSrS(i) = variant.LogitsTransform(mainloop_params, tSrS(i), /*batch_idx=*/0, qo_idx, kv_idx, + qo_head_idx, kv_head_idx); + if constexpr (MULTIITEMSCORING) { + mask_multi_item_scoring(tSrS, i, qo_idx, kv_idx); + } else if constexpr (!CAUSAL) { // Just masking based on col + if (kv_idx >= kv_len) { + tSrS(i) = AttentionUpdater::fill_value; + } + } else { + if (kv_idx >= std::min(kv_len, col_limit_right(qo_idx))) { + tSrS(i) = AttentionUpdater::fill_value; + } + } + if constexpr (LEFT_SLIDING_WINDOW) { + if (kv_idx < col_limit_left(qo_idx)) { + tSrS(i) = AttentionUpdater::fill_value; + } + } + } + } + + attention_updater.update(tSrS); + Tensor tOrP = make_tensor(convert_type(tSrS).data(), + convert_layout_acc_Aregs(tSrS.layout())); + + constexpr int n_masking_steps = MULTIITEMSCORING ? (cute::ceil_div(CTA_Q, CTA_KV) + 1) + : (CAUSAL ? cute::ceil_div(CTA_Q, CTA_KV) : 0); + // masking loops + // ziangl@nvidia.com: for multi item scoring, we use this loop only to mask along the diagonal +#pragma unroll + for (int masking_step = 0; masking_step < n_masking_steps && kv_tile_idx > swa_begin_kv_tile_idx; + ++masking_step, kv_tile_idx = kv_tile_idx_decrement(kv_tile_idx)) { + Tensor tSrS = partition_fragment_C(tiled_mma_qk, select<0, 1>(TileShape_QKD{})); + consumer_wait(pipeline_k, smem_pipe_read_k); + WarpScheduler::barrier_sync(); + gemm(tiled_mma_qk, tSrQ, tSrK(_, _, _, smem_pipe_read_k.index()), + tSrS); + if (masking_step > 0) { + attention_updater.rescale_o(tOrO); + } + consumer_wait(pipeline_v, smem_pipe_read_v); + gemm(tiled_mma_pv, tOrP, + tOrV(_, _, _, smem_pipe_read_v.index()), tOrO); + WarpScheduler::barrier_arrive(); + warpgroup_wait<1>(); + pipeline_k.consumer_release(smem_pipe_read_k); // release K + Tensor cS = cute::make_identity_tensor(select<0, 1>(TileShape_QKD{})); + Tensor tScS = threadMmaQK.partition_C(cS); +#pragma unroll + for (int i = 0; i < size(tSrS); ++i) { + int qo_idx = get<0>(tScS(i)) + q_tile_idx * CTA_Q; + int kv_idx = get<1>(tScS(i)) + kv_tile_idx_decrement(kv_tile_idx) * CTA_KV; + tSrS(i) = variant.LogitsTransform(mainloop_params, tSrS(i), /*batch_idx=*/0, qo_idx, kv_idx, + qo_head_idx, kv_head_idx); + if (MULTIITEMSCORING) { + mask_multi_item_scoring(tSrS, i, qo_idx, kv_idx); + } else { + if (kv_idx >= col_limit_right(qo_idx)) { + tSrS(i) = AttentionUpdater::fill_value; + } + } + if constexpr (LEFT_SLIDING_WINDOW) { + if (kv_idx < col_limit_left(qo_idx)) { + tSrS(i) = AttentionUpdater::fill_value; + } + } + } + attention_updater.update(tSrS); + warpgroup_wait<0>(); + pipeline_v.consumer_release(smem_pipe_read_v); // release V + ++smem_pipe_read_k; + ++smem_pipe_read_v; + cute::copy(make_tensor(convert_type(tSrS).data(), + convert_layout_acc_Aregs(tSrS.layout())), + tOrP); + } + +#pragma unroll 1 + for (; kv_tile_idx > swa_end_kv_tile_idx + 1; kv_tile_idx = kv_tile_idx_decrement(kv_tile_idx)) { + Tensor tSrS = partition_fragment_C(tiled_mma_qk, select<0, 1>(TileShape_QKD{})); + consumer_wait(pipeline_k, smem_pipe_read_k); + WarpScheduler::barrier_sync(); + gemm(tiled_mma_qk, tSrQ, tSrK(_, _, _, smem_pipe_read_k.index()), + tSrS); + attention_updater.rescale_o(tOrO); + consumer_wait(pipeline_v, smem_pipe_read_v); + gemm(tiled_mma_pv, tOrP, + tOrV(_, _, _, smem_pipe_read_v.index()), tOrO); + WarpScheduler::barrier_arrive(); + warpgroup_wait<1>(); + pipeline_k.consumer_release(smem_pipe_read_k); // release K + // #pragma unroll + Tensor cS = cute::make_identity_tensor(select<0, 1>(TileShape_QKD{})); + Tensor tScS = threadMmaQK.partition_C(cS); +#pragma unroll + for (int i = 0; i < size(tSrS); ++i) { + int qo_idx = get<0>(tScS(i)) + q_tile_idx * CTA_Q; + int kv_idx = get<1>(tScS(i)) + kv_tile_idx_decrement(kv_tile_idx) * CTA_KV; + tSrS(i) = variant.LogitsTransform(mainloop_params, tSrS(i), /*batch_idx=*/0, qo_idx, kv_idx, + qo_head_idx, kv_head_idx); + } + if constexpr (MULTIITEMSCORING) { + // auto nums_tiles_outside_causal_diagonal = kv_tile_idx_count - cute::ceil_div(CTA_Q, + // CTA_KV); + if (kv_tile_idx >= num_kv_tiles_prefix - 1) { +#pragma unroll + for (int i = 0; i < size(tSrS); ++i) { + int qo_idx = get<0>(tScS(i)) + q_tile_idx * CTA_Q; + int kv_idx = get<1>(tScS(i)) + kv_tile_idx_decrement(kv_tile_idx) * CTA_KV; + mask_multi_item_scoring_assume_in_bound(tSrS, i, qo_idx, kv_idx); + } + } + } + attention_updater.update(tSrS); + warpgroup_wait<0>(); + pipeline_v.consumer_release(smem_pipe_read_v); // release V + ++smem_pipe_read_k; + ++smem_pipe_read_v; + cute::copy(make_tensor(convert_type(tSrS).data(), + convert_layout_acc_Aregs(tSrS.layout())), + tOrP); + } + + if constexpr (LEFT_SLIDING_WINDOW) { +#pragma unroll 1 + for (; kv_tile_idx > swa_begin_kv_tile_idx; --kv_tile_idx) { + Tensor tSrS = partition_fragment_C(tiled_mma_qk, select<0, 1>(TileShape_QKD{})); + consumer_wait(pipeline_k, smem_pipe_read_k); + WarpScheduler::barrier_sync(); + gemm(tiled_mma_qk, tSrQ, + tSrK(_, _, _, smem_pipe_read_k.index()), tSrS); + attention_updater.rescale_o(tOrO); + consumer_wait(pipeline_v, smem_pipe_read_v); + gemm(tiled_mma_pv, tOrP, + tOrV(_, _, _, smem_pipe_read_v.index()), tOrO); + WarpScheduler::barrier_arrive(); + warpgroup_wait<1>(); + pipeline_k.consumer_release(smem_pipe_read_k); // release K + Tensor cS = cute::make_identity_tensor(select<0, 1>(TileShape_QKD{})); + Tensor tScS = threadMmaQK.partition_C(cS); +#pragma unroll + for (int i = 0; i < size(tSrS); ++i) { + int qo_idx = get<0>(tScS(i)) + q_tile_idx * CTA_Q; + int kv_idx = get<1>(tScS(i)) + (kv_tile_idx - 1) * CTA_KV; + tSrS(i) = variant.LogitsTransform(mainloop_params, tSrS(i), /*batch_idx=*/0, qo_idx, kv_idx, + qo_head_idx, kv_head_idx); + if (kv_idx < col_limit_left(qo_idx)) { + tSrS(i) = AttentionUpdater::fill_value; + } + } + attention_updater.update(tSrS); + warpgroup_wait<0>(); + pipeline_v.consumer_release(smem_pipe_read_v); // release V + ++smem_pipe_read_k; + ++smem_pipe_read_v; + cute::copy(make_tensor(convert_type(tSrS).data(), + convert_layout_acc_Aregs(tSrS.layout())), + tOrP); + } + } + + // Tell warp 0 that smem_q is ready + cutlass::arch::NamedBarrier::arrive(NUM_MMA_THREADS + Ktraits::NUM_PRODUCER_THREADS, + /*id=*/static_cast(NamedBarriers::kQueryEmpty)); + attention_updater.rescale_o(tOrO); + consumer_wait(pipeline_v, smem_pipe_read_v); + gemm(tiled_mma_pv, tOrP, tOrV(_, _, _, smem_pipe_read_v.index()), + tOrO); + attention_updater.finalize(tSrS); + warpgroup_wait<0>(); + pipeline_v.consumer_release(smem_pipe_read_v); // release V, otherwise producers will hang + ++smem_pipe_read_v; + + attention_updater.rescale_o(tOrO); + return; +} + +} // namespace flashinfer + +#endif // FLASHINFER_ATTENTION_HOPPER_MAINLOOP_MMA_CUH_ diff --git a/include/flashinfer/attention/hopper/named_barrier.cuh b/include/flashinfer/attention/hopper/named_barrier.cuh new file mode 100644 index 0000000000000000000000000000000000000000..b28dd92eb5bb38f4f90acf0a6d7c066b127615c6 --- /dev/null +++ b/include/flashinfer/attention/hopper/named_barrier.cuh @@ -0,0 +1,113 @@ +/* + * Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri + * Dao. Licensed under the BSD 3-Clause. + * + * Modified by the FlashInfer team. + */ +#ifndef FLASHINFER_ATTENTION_HOPPER_NAMED_BARRIERS_CUH_ +#define FLASHINFER_ATTENTION_HOPPER_NAMED_BARRIERS_CUH_ + +#include + +#include "cutlass/arch/barrier.h" +#include "cutlass/cutlass.h" + +namespace flashinfer { + +// Enumerates the reserved named barriers to avoid potential conflicts + +enum class NamedBarriers { + kQueryEmpty = 0, + kValueEmpty = 1, + kWarpSchedulerWG1 = 2, + kWarpSchedulerWG2 = 3, + kWarpSchedulerWG3 = 4, + kPrefetchIndices = 5, + kProducerWG = 6 +}; + +__device__ __forceinline__ int get_warp_group_barrier_idx(int warp_group_idx) { + return static_cast(NamedBarriers::kWarpSchedulerWG1) + warp_group_idx - 1; +} + +template +__device__ __forceinline__ int get_next_consumer_warp_group_idx() { + static_assert(num_consumer_warp_groups == 2 || num_consumer_warp_groups == 3); + int warp_group_idx = cutlass::canonical_warp_group_idx(); + if constexpr (num_consumer_warp_groups == 2) { + // 1 -> 2, 2 -> 1 + return 3 - warp_group_idx; + } else { + // num_consumer_warp_groups == 3 + // 1 -> 2, 2 -> 3, 3 -> 1 + return (warp_group_idx % 3) + 1; + } +} + +template +__device__ __forceinline__ int get_prev_consumer_warp_group_idx() { + static_assert(num_consumer_warp_groups == 2 || num_consumer_warp_groups == 3); + int warp_group_idx = cutlass::canonical_warp_group_idx(); + if constexpr (num_consumer_warp_groups == 2) { + // 1 -> 2, 2 -> 1 + return 3 - warp_group_idx; + } else { + // num_consumer_warp_groups == 3 + // 1 -> 3, 2 -> 1, 3 -> 2 + return ((warp_group_idx + 1) % 3) + 1; + } +} + +template +struct WarpScheduler { + constexpr static int NUM_MMA_THREADS = Ktraits::NUM_MMA_THREADS; + static CUTLASS_DEVICE void barrier_sync() { + if constexpr (UseSchedulerBarrier) { + cutlass::arch::NamedBarrier::sync( + NUM_MMA_THREADS, get_warp_group_barrier_idx(cutlass::canonical_warp_group_idx())); + } + } + + static CUTLASS_DEVICE void barrier_arrive() { + if constexpr (!UseSchedulerBarrier) { + return; + } + static_assert(NUM_MMA_THREADS == 2 * cutlass::NumThreadsPerWarpGroup || + NUM_MMA_THREADS == 3 * cutlass::NumThreadsPerWarpGroup); + if constexpr (NUM_MMA_THREADS == 2 * cutlass::NumThreadsPerWarpGroup) { + cutlass::arch::NamedBarrier::arrive( + NUM_MMA_THREADS, get_warp_group_barrier_idx(get_next_consumer_warp_group_idx<2>())); + } else { + cutlass::arch::NamedBarrier::arrive( + NUM_MMA_THREADS, get_warp_group_barrier_idx(get_next_consumer_warp_group_idx<3>())); + cutlass::arch::NamedBarrier::arrive( + NUM_MMA_THREADS, get_warp_group_barrier_idx(get_prev_consumer_warp_group_idx<3>())); + } + } + + static CUTLASS_DEVICE void mma_init() { + // Tell producer (warp 0) that smem_q is ready + cutlass::arch::NamedBarrier::arrive(NUM_MMA_THREADS + Ktraits::NUM_PRODUCER_THREADS, + /*id=*/static_cast(NamedBarriers::kQueryEmpty)); + if constexpr (!UseSchedulerBarrier) { + return; + } + static_assert(NUM_MMA_THREADS == 2 * cutlass::NumThreadsPerWarpGroup || + NUM_MMA_THREADS == 3 * cutlass::NumThreadsPerWarpGroup); + if (cutlass::canonical_warp_group_idx() > 1) { + cutlass::arch::NamedBarrier::arrive( + NUM_MMA_THREADS, /*id=*/static_cast(NamedBarriers::kWarpSchedulerWG1)); + } + if constexpr (NUM_MMA_THREADS == 3 * cutlass::NumThreadsPerWarpGroup) { + if (cutlass::canonical_warp_group_idx() > 2) { + cutlass::arch::NamedBarrier::arrive( + NUM_MMA_THREADS, /*id=*/static_cast(NamedBarriers::kWarpSchedulerWG2)); + } + } + } + +}; // struct WarpScheduler + +} // namespace flashinfer + +#endif // FLASHINFER_ATTENTION_HOPPER_NAMED_BARRIERS_CUH_ diff --git a/include/flashinfer/attention/hopper/prefill_sm90.cuh b/include/flashinfer/attention/hopper/prefill_sm90.cuh new file mode 100644 index 0000000000000000000000000000000000000000..da7fda25be7e1b1c94fcc197c571885f8b41bbb4 --- /dev/null +++ b/include/flashinfer/attention/hopper/prefill_sm90.cuh @@ -0,0 +1,618 @@ +/* + * Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri + * Dao. Licensed under the BSD 3-Clause. + * + * Modified by the FlashInfer team. + */ +#ifndef FLASHINFER_ATTENTION_HOPPER_PREFILL_SM90_CUH_ +#define FLASHINFER_ATTENTION_HOPPER_PREFILL_SM90_CUH_ + +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +#include "../../cutlass_utils.cuh" +#include "../../exception.h" +#include "../mask.cuh" +#include "cute/tensor.hpp" +#include "cutlass/pipeline/pipeline.hpp" +#include "epilogue.cuh" +#include "kernel_traits.cuh" +#include "mainloop.cuh" +#include "mainloop_mma.cuh" +#include "sparse_mainloop.cuh" +#include "tile_scheduler.cuh" +#include "utils.cuh" + +namespace flashinfer { + +using namespace cute; + +DEFINE_HAS_MEMBER(maybe_prefix_len_ptr) +DEFINE_HAS_MEMBER(maybe_token_pos_in_items_ptr) +DEFINE_HAS_MEMBER(token_pos_in_items_len) +DEFINE_HAS_MEMBER(maybe_max_item_len_ptr) + +template +__global__ void __launch_bounds__(Ktraits::NUM_WARPS* cutlass::NumThreadsPerWarp, 1) + PrefillWithKVCacheKernel(CUTE_GRID_CONSTANT + typename CollectiveMainloop::Params const mainloop_params, + CUTE_GRID_CONSTANT + typename CollectiveEpilogue::Params const epilogue_params, + CUTE_GRID_CONSTANT + typename TileScheduler::Params const scheduler_params) { + using DTypeQ = typename Ktraits::DTypeQ; + using DTypeKV = typename Ktraits::DTypeKV; + using DTypeO = typename Ktraits::DTypeO; + using DTypeQKAccum = typename Ktraits::DTypeQKAccum; + using TileShape_QKD = typename Ktraits::TileShape_QKD; + using TileShape_PDV = typename Ktraits::TileShape_PDV; + using AttentionVariant = typename Ktraits::AttentionVariant; + + static constexpr int NUM_MMA_THREADS = Ktraits::NUM_MMA_THREADS; + static constexpr int NUM_COPY_THREADS = cutlass::NumThreadsPerWarpGroup; + static constexpr int CTA_Q = Ktraits::CTA_Q; + static constexpr int CTA_KV = Ktraits::CTA_KV; + + static constexpr bool use_tma_load_kv = CollectiveMainloop::USE_TMA_LOAD_KV; + + using MainloopPipeline = typename CollectiveMainloop::MainloopPipeline; + using PipelineParams = typename MainloopPipeline::Params; + using PipelineState = typename MainloopPipeline::PipelineState; + + extern __shared__ char shared_memory[]; + auto& shared_storage = *reinterpret_cast(shared_memory); + + int const lane_predicate = cute::elect_one_sync(); + int const warp_idx = cutlass::canonical_warp_idx_sync(); + + // Issue Tma Descriptor Prefetch from a single thread + if (warp_idx == 0 && lane_predicate) { + CollectiveMainloop::prefetch_tma_descriptors(mainloop_params); + CollectiveEpilogue::prefetch_tma_descriptors(epilogue_params); + } + + // Obtain warp index + int const warp_group_thread_idx = threadIdx.x % cutlass::NumThreadsPerWarpGroup; + + PipelineParams pipeline_params; + int warp_group_idx = cutlass::canonical_warp_group_idx(); + pipeline_params.role = warp_group_idx == 0 ? MainloopPipeline::ThreadCategory::Producer + : MainloopPipeline::ThreadCategory::Consumer; + if constexpr (use_tma_load_kv) { + pipeline_params.is_leader = warp_group_thread_idx == 0; + pipeline_params.num_consumers = NUM_MMA_THREADS; + } else { + pipeline_params.producer_arv_count = NUM_COPY_THREADS; + pipeline_params.consumer_arv_count = NUM_MMA_THREADS; + } + + if (warp_idx == 0 && lane_predicate) { + shared_storage.barrier_Q.init(/*num_threads=*/1); + shared_storage.barrier_O.init(/*num_threads=*/1); + } + // We're counting on pipeline_k to call cutlass::arch::fence_barrier_init(); + MainloopPipeline pipeline_k = [&] { + if constexpr (use_tma_load_kv) { + pipeline_params.transaction_bytes = CollectiveMainloop::TmaTransactionBytesK; + return MainloopPipeline(shared_storage.pipeline_k, pipeline_params, + /*cluster_shape=*/Shape<_1, _1, _1>{}); + } else { + return MainloopPipeline(shared_storage.pipeline_k, pipeline_params); + } + }(); + + MainloopPipeline pipeline_v = [&] { + if constexpr (use_tma_load_kv) { + pipeline_params.transaction_bytes = CollectiveMainloop::TmaTransactionBytesV; + return MainloopPipeline(shared_storage.pipeline_v, pipeline_params, + /*cluster_shape=*/Shape<_1, _1, _1>{}); + } else { + return MainloopPipeline(shared_storage.pipeline_v, pipeline_params); + } + }(); + + CollectiveMainloop collective_mainloop; + CollectiveEpilogue collective_epilogue; + + // We need this to guarantee that the Pipeline init is visible to all producers and consumer + // blocks in the Cluster + __syncthreads(); + + uint32_t* maybe_prefix_len_ptr = nullptr; + if constexpr (has_maybe_prefix_len_ptr_v) { + maybe_prefix_len_ptr = mainloop_params.additional_params.maybe_prefix_len_ptr; + } + uint16_t* maybe_token_pos_in_items_ptr = nullptr; + if constexpr (has_maybe_token_pos_in_items_ptr_v) { + maybe_token_pos_in_items_ptr = mainloop_params.additional_params.maybe_token_pos_in_items_ptr; + } + uint32_t token_pos_in_items_len = 0; + if constexpr (has_token_pos_in_items_len_v) { + token_pos_in_items_len = mainloop_params.additional_params.token_pos_in_items_len; + } + uint16_t* maybe_max_item_len_ptr = nullptr; + if constexpr (has_maybe_max_item_len_ptr_v) { + maybe_max_item_len_ptr = mainloop_params.additional_params.maybe_max_item_len_ptr; + } + + if (warp_group_idx == 0) { // Producer + if constexpr (use_tma_load_kv) { + cutlass::arch::warpgroup_reg_dealloc(); + } else { + cutlass::arch::warpgroup_reg_dealloc<72>(); + } + + int warp_idx_in_warpgroup = __shfl_sync(0xffffffff, (threadIdx.x / 32) % 4, 0); + if (!use_tma_load_kv || warp_idx_in_warpgroup == 0) { // Load Q, K, V + PipelineState smem_pipe_write_k = cutlass::make_producer_start_state(); + PipelineState smem_pipe_write_v = cutlass::make_producer_start_state(); + + int work_idx = 0; + + TileScheduler scheduler; + for (auto work_tile_info = scheduler.get_initial_work(scheduler_params); + work_tile_info.is_valid(scheduler_params); + work_tile_info = scheduler.template get_next_work( + scheduler_params, work_tile_info)) { + auto block_coord = work_tile_info.get_block_coord(scheduler_params); + auto [q_tile_idx, qo_head_idx, kv_head_idx, qo_indptr, kv_indptr, qo_len, kv_len, + batch_idx] = block_coord; + + if (q_tile_idx * CTA_Q >= qo_len) { + continue; + } + int num_kv_tiles = + collective_mainloop.get_num_kv_tiles(mainloop_params, q_tile_idx, qo_len, kv_len); + if (num_kv_tiles <= 0) { + scheduler.prefetch_next_work(scheduler_params, work_tile_info); + scheduler.broadcast_next_work(work_tile_info); + continue; + } + int num_kv_tiles_outside_items_window = 0; + int num_kv_tiles_prefix = 0; + if constexpr (MULTIITEMSCORING) { + auto prefix_len = __ldg(maybe_prefix_len_ptr + batch_idx); + auto max_item_len = __ldg(maybe_max_item_len_ptr + batch_idx); + auto valid_items_window_len = + std::max(0, q_tile_idx * CTA_Q + kv_len - qo_len - max_item_len); + num_kv_tiles_outside_items_window = valid_items_window_len / CTA_KV; + num_kv_tiles_prefix = cute::ceil_div(prefix_len, CTA_KV); + } + if constexpr (MULTIITEMSCORING) { + collective_mainloop.load( + mainloop_params, pipeline_k, pipeline_v, smem_pipe_write_k, smem_pipe_write_v, + shared_storage, scheduler, scheduler_params, work_tile_info, block_coord, work_idx, + num_kv_tiles_outside_items_window, num_kv_tiles_prefix); + } else { + collective_mainloop.load( + mainloop_params, pipeline_k, pipeline_v, smem_pipe_write_k, smem_pipe_write_v, + shared_storage, scheduler, scheduler_params, work_tile_info, block_coord, work_idx); + } + ++work_idx; + } + collective_mainloop.load_tail(pipeline_k, pipeline_v, smem_pipe_write_k, smem_pipe_write_v); + } + } else { // Consumer + if constexpr (use_tma_load_kv) { + cutlass::arch::warpgroup_reg_alloc(); + } else { + cutlass::arch::warpgroup_reg_alloc(); + } + + TileScheduler scheduler; + // Initialize matmul objects. + typename Ktraits::TiledMmaPV tiled_mma_pv; + + PipelineState smem_pipe_read_k, smem_pipe_read_v; + // We don't need separate variables smem_pipe_release_k and smem_pipe_release_v + // (like in Cutlass's gemm) because the read and release pipeline states are always the same. + + CollectiveMainloop::WarpScheduler::mma_init(); + scheduler.init_consumer(); + + int work_idx = 0; + CUTLASS_PRAGMA_NO_UNROLL + for (auto work_tile_info = scheduler.get_initial_work(scheduler_params); + work_tile_info.is_valid(scheduler_params); + work_tile_info = scheduler.template get_next_work(scheduler_params, + work_tile_info)) { + // Attention output (GEMM-II) accumulator. + Tensor tOrO = partition_fragment_C(tiled_mma_pv, select<0, 1>(TileShape_PDV{})); + + auto block_coord = work_tile_info.get_block_coord(scheduler_params); + auto [q_tile_idx, qo_head_idx, kv_head_idx, qo_indptr, kv_indptr, qo_len, kv_len, batch_idx] = + block_coord; + + AttentionVariant variant(mainloop_params, block_coord); + auto attention_updater = + variant.template GetAttentionUpdater<2 * (2 * CTA_Q / NUM_MMA_THREADS)>(); + + if (q_tile_idx * CTA_Q >= qo_len) { + continue; + } + int num_kv_tiles = + collective_mainloop.get_num_kv_tiles(mainloop_params, q_tile_idx, qo_len, kv_len); + if (num_kv_tiles <= 0) { // We exit early and write 0 to gO and -inf to gLSE. + collective_epilogue.store_zero(epilogue_params, shared_storage, + threadIdx.x - NUM_COPY_THREADS, block_coord); + continue; + } + + int swa_begin_kv_tile_idx = 0; + int swa_end_kv_tile_idx = -1; + if constexpr (LEFT_SLIDING_WINDOW) { + swa_begin_kv_tile_idx = get_swa_begin_kv_tile_idx( + mainloop_params.window_left, q_tile_idx, qo_len, kv_len); + swa_end_kv_tile_idx = get_swa_end_kv_tile_idx(mainloop_params.window_left, + q_tile_idx, qo_len, kv_len); + } + + uint32_t prefix_len = 0; + uint16_t* token_pos_in_items = nullptr; + if constexpr (MULTIITEMSCORING) { + prefix_len = __ldg(maybe_prefix_len_ptr + batch_idx); + token_pos_in_items = maybe_token_pos_in_items_ptr + batch_idx * token_pos_in_items_len; + } + int num_kv_tiles_outside_items_window = 0; + int num_kv_tiles_prefix = 0; + if constexpr (MULTIITEMSCORING) { + auto prefix_len = __ldg(maybe_prefix_len_ptr + batch_idx); + auto max_item_len = __ldg(maybe_max_item_len_ptr + batch_idx); + auto valid_items_window_len = + std::max(0, q_tile_idx * CTA_Q + kv_len - qo_len - max_item_len); + num_kv_tiles_outside_items_window = valid_items_window_len / CTA_KV; + num_kv_tiles_prefix = cute::ceil_div(prefix_len, CTA_KV); + } + mma_f16( + mainloop_params, variant, pipeline_k, pipeline_v, smem_pipe_read_k, smem_pipe_read_v, + tOrO, attention_updater, num_kv_tiles, swa_begin_kv_tile_idx, swa_end_kv_tile_idx, + threadIdx.x - NUM_COPY_THREADS, work_idx, q_tile_idx, shared_storage, qo_len, kv_len, + qo_head_idx, kv_head_idx, prefix_len, token_pos_in_items, + num_kv_tiles_outside_items_window, num_kv_tiles_prefix); + collective_epilogue.store(epilogue_params, tOrO, attention_updater.get_lse(), shared_storage, + tiled_mma_pv, threadIdx.x - NUM_COPY_THREADS, block_coord); + + ++work_idx; + } + collective_epilogue.store_tail(); + } +} + +template +cudaError_t SinglePrefillWithKVCacheKernelTraitsDispatched(Params& params, cudaStream_t stream) { + using DTypeQ = typename KernelTraits::DTypeQ; + using DTypeKV = typename KernelTraits::DTypeKV; + using DTypeO = typename KernelTraits::DTypeO; + + using CollectiveMainloop = + CollectiveMainloop; + using CollectiveEpilogue = CollectiveEpilogue; + using Scheduler = SingleTileScheduler; + typename CollectiveMainloop::Params mainloop_params = CollectiveMainloop::to_underlying_arguments( + {params.q_ptr, + get_gmem_layout(params.qo_len, params.num_qo_heads, KernelTraits::HEAD_DIM_QK, + params.q_stride_n, + params.q_stride_h), // layout_Q + params.k_ptr, + get_gmem_layout(params.kv_len, params.num_kv_heads, KernelTraits::HEAD_DIM_QK, + params.k_stride_n, + params.k_stride_h), // layout_K + params.v_ptr, + get_gmem_layout(params.kv_len, params.num_kv_heads, KernelTraits::HEAD_DIM_VO, + params.v_stride_n, + params.v_stride_h), // layout_V + params.window_left, params.additional_params}); + typename CollectiveEpilogue::Params epilogue_params = + CollectiveEpilogue::to_underlying_arguments({ + static_cast(params.o_ptr), + get_gmem_layout(params.qo_len, params.num_qo_heads, KernelTraits::HEAD_DIM_VO, + params.o_stride_n, + params.o_stride_h), // layout_O + static_cast(params.lse_ptr), + get_lse_gmem_layout(params.qo_len, params.num_qo_heads), // layout_LSE + }); + + int num_tiles_q = cutlass::ceil_div(params.qo_len, KernelTraits::CTA_Q); + // TODO(Zihao): also support kv-head major + typename Scheduler::Arguments scheduler_args = { + num_tiles_q, params.num_qo_heads, params.qo_len, params.kv_len, + cutlass::FastDivmod(params.num_qo_heads / params.num_kv_heads)}; + typename Scheduler::Params scheduler_params = Scheduler::to_underlying_arguments(scheduler_args); + + auto kernel = + (void*)PrefillWithKVCacheKernel; + int smem_size = sizeof(typename KernelTraits::SharedStorage); + FLASHINFER_CUDA_CALL( + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + + int device; + cudaGetDevice(&device); + int multiprocessor_count; + FLASHINFER_CUDA_CALL( + cudaDeviceGetAttribute(&multiprocessor_count, cudaDevAttrMultiProcessorCount, device)); + dim3 grid_dims = Scheduler::get_grid_dim(scheduler_args, multiprocessor_count); + static constexpr int num_ctas = KernelTraits::NUM_WARPS * 32; + dim3 block_dims(num_ctas); + void* args[] = {&mainloop_params, &epilogue_params, &scheduler_params}; + FLASHINFER_CUDA_CALL(cudaLaunchKernel(kernel, grid_dims, block_dims, args, smem_size, stream)); + + return cudaSuccess; +} + +template +cudaError_t BatchPrefillWithPagedKVCacheKernelTraitsDispatched(Params& params, + cudaStream_t stream) { + using DTypeQ = typename KernelTraits::DTypeQ; + using DTypeKV = typename KernelTraits::DTypeKV; + using DTypeO = typename KernelTraits::DTypeO; + using IdType = typename KernelTraits::IdType; + + using CollectiveMainloop = SparseCollectiveMainloop; + using CollectiveEpilogue = CollectiveEpilogue; + using Scheduler = + std::conditional_t, + BatchPrefillPersistentTileScheduler>; + + typename CollectiveMainloop::Params mainloop_params = CollectiveMainloop::to_underlying_arguments( + {params.q_ptr, + get_gmem_layout(params.nnz_qo, params.num_qo_heads, KernelTraits::HEAD_DIM_QK, + params.q_stride_n, + params.q_stride_h), // layout_Q + params.k_ptr, + // NOTE(Zihao): nnz was useless here, we can just pass 0 + get_gmem_layout(/*nnz=*/0, params.num_kv_heads, KernelTraits::HEAD_DIM_QK, params.k_stride_n, + params.k_stride_h), // layout_K + params.v_ptr, + get_gmem_layout(/*nnz=*/0, params.num_kv_heads, KernelTraits::HEAD_DIM_VO, params.v_stride_n, + params.v_stride_h), // layout_V + params.kv_indices, params.window_left, params.additional_params}); + typename CollectiveEpilogue::Params epilogue_params = + CollectiveEpilogue::to_underlying_arguments({ + params.o_ptr, + get_gmem_layout(params.nnz_qo, params.num_qo_heads, KernelTraits::HEAD_DIM_VO, + params.o_stride_n, + params.o_stride_h), // layout_O + params.lse_ptr, get_lse_gmem_layout(params.nnz_qo, params.num_qo_heads), // layout_LSE + }); + + typename Scheduler::Arguments scheduler_args = { + params.work_indptr, + params.head_indices, + params.qo_tile_indices, + params.qo_indptr, + params.kv_indptr, + params.qo_lens, + params.kv_lens, + params.batch_indices, + cutlass::FastDivmod(params.num_qo_heads / params.num_kv_heads), + params.num_qo_heads}; + typename Scheduler::Params scheduler_params = Scheduler::to_underlying_arguments(scheduler_args); + + // Get the ptr to kernel function. + auto kernel = + (void*)PrefillWithKVCacheKernel; + int smem_size = sizeof(typename KernelTraits::SharedStorage); + FLASHINFER_CUDA_CALL( + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + + int device; + cudaGetDevice(&device); + int multiprocessor_count; + FLASHINFER_CUDA_CALL( + cudaDeviceGetAttribute(&multiprocessor_count, cudaDevAttrMultiProcessorCount, device)); + dim3 grid_dims = Scheduler::get_grid_dim(scheduler_args, multiprocessor_count); + static constexpr int ctaSize = KernelTraits::NUM_WARPS * 32; + dim3 block_dims(ctaSize); + void* args[] = {&mainloop_params, &epilogue_params, &scheduler_params}; + FLASHINFER_CUDA_CALL(cudaLaunchKernel(kernel, grid_dims, block_dims, args, smem_size, stream)); + + return cudaSuccess; +} + +template +cudaError_t BatchPrefillWithRaggedKVCacheKernelTraitsDispatched(Params& params, + cudaStream_t stream) { + using DTypeQ = typename KernelTraits::DTypeQ; + using DTypeKV = typename KernelTraits::DTypeKV; + using DTypeO = typename KernelTraits::DTypeO; + using IdType = typename KernelTraits::IdType; + + using CollectiveMainloop = + CollectiveMainloop; + using CollectiveEpilogue = CollectiveEpilogue; + using Scheduler = + std::conditional_t, + BatchPrefillPersistentTileScheduler>; + typename CollectiveMainloop::Params mainloop_params = CollectiveMainloop::to_underlying_arguments( + {params.q_ptr, + get_gmem_layout(params.nnz_qo, params.num_qo_heads, KernelTraits::HEAD_DIM_QK, + params.q_stride_n, + params.q_stride_h), // layout_Q + params.k_ptr, + // NOTE(Zihao): nnz was useless here, we can just pass 0 + get_gmem_layout(params.nnz_kv, params.num_kv_heads, KernelTraits::HEAD_DIM_QK, + params.k_stride_n, + params.k_stride_h), // layout_K + params.v_ptr, + get_gmem_layout(params.nnz_kv, params.num_kv_heads, KernelTraits::HEAD_DIM_VO, + params.v_stride_n, + params.v_stride_h), // layout_V + params.window_left, params.additional_params}); + typename CollectiveEpilogue::Params epilogue_params = + CollectiveEpilogue::to_underlying_arguments({ + params.o_ptr, + get_gmem_layout(params.nnz_qo, params.num_qo_heads, KernelTraits::HEAD_DIM_VO, + params.o_stride_n, + params.o_stride_h), // layout_O + params.lse_ptr, get_lse_gmem_layout(params.nnz_qo, params.num_qo_heads), // layout_LSE + }); + + // NOTE(Zihao): add support for kv head-major later + typename Scheduler::Arguments scheduler_args = { + params.work_indptr, + params.head_indices, + params.qo_tile_indices, + params.qo_indptr, + params.kv_indptr, + params.qo_lens, + params.kv_lens, + params.batch_indices, + cutlass::FastDivmod(params.num_qo_heads / params.num_kv_heads), + params.num_qo_heads}; + typename Scheduler::Params scheduler_params = Scheduler::to_underlying_arguments(scheduler_args); + + // Get the ptr to kernel function. + auto kernel = + (void*)PrefillWithKVCacheKernel; + int smem_size = sizeof(typename KernelTraits::SharedStorage); + FLASHINFER_CUDA_CALL( + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + + int device; + cudaGetDevice(&device); + int multiprocessor_count; + FLASHINFER_CUDA_CALL( + cudaDeviceGetAttribute(&multiprocessor_count, cudaDevAttrMultiProcessorCount, device)); + dim3 grid_dims = Scheduler::get_grid_dim(scheduler_args, multiprocessor_count); + static constexpr int ctaSize = KernelTraits::NUM_WARPS * 32; + dim3 block_dims(ctaSize); + void* args[] = {&mainloop_params, &epilogue_params, &scheduler_params}; + FLASHINFER_CUDA_CALL(cudaLaunchKernel(kernel, grid_dims, block_dims, args, smem_size, stream)); + + return cudaSuccess; +} + +template +constexpr auto getCTATileSize() { + if constexpr (HEAD_DIM_QK == HEAD_DIM_VO) { + if constexpr (HEAD_DIM_QK == 64) { + return std::make_tuple(192, 128); + } else if constexpr (HEAD_DIM_QK == 128) { + if constexpr (CAUSAL) { + return std::make_tuple(128, 128); + } else { + return std::make_tuple(128, 192); + } + } else { + return std::make_tuple(128, 64); + } + } else { + // NOTE(Zihao) hack for deepseek prefill + static_assert(HEAD_DIM_QK == 192 && HEAD_DIM_VO == 128); + return std::make_tuple(128, 128); + } +} + +template +cudaError_t SinglePrefillWithKVCacheDispatched(Params& params, cudaStream_t stream) { + static_assert(HEAD_DIM_VO == 64 || HEAD_DIM_VO == 128 || HEAD_DIM_VO == 256); + if (MASK_MODE == MaskMode::kCustom) { + return cudaErrorNotSupported; // Not supported yet. + } + constexpr bool CAUSAL = MASK_MODE == MaskMode::kCausal; + constexpr auto CTA_TILE_SIZE = getCTATileSize(); + SinglePrefillWithKVCacheKernelTraitsDispatched< + AttentionKernelTraits(CTA_TILE_SIZE), + /*CTA_KV_=*/get<1>(CTA_TILE_SIZE), + /*NUM_STAGES_=*/2, typename Params::DTypeQ, typename Params::DTypeKV, + typename Params::DTypeO, typename Params::IdType, AttentionVariant>, + LEFT_SLIDING_WINDOW, CAUSAL>(params, stream); + cudaError_t status = cudaGetLastError(); + return status; +} + +template +cudaError_t BatchPrefillWithRaggedKVCacheDispatched(Params& params, bool enable_pdl, + cudaStream_t stream) { + static_assert(HEAD_DIM_VO == 64 || HEAD_DIM_VO == 128 || HEAD_DIM_VO == 256); + if (MASK_MODE == MaskMode::kCustom) { + return cudaErrorNotSupported; // Not supported yet. + } + constexpr bool CAUSAL = MASK_MODE == MaskMode::kCausal; + constexpr auto CTA_TILE_SIZE = getCTATileSize(); + BatchPrefillWithRaggedKVCacheKernelTraitsDispatched< + AttentionKernelTraits(CTA_TILE_SIZE), + /*CTA_KV_=*/get<1>(CTA_TILE_SIZE), + /*NUM_STAGES_=*/2, typename Params::DTypeQ, typename Params::DTypeKV, + typename Params::DTypeO, typename Params::IdType, AttentionVariant>, + LEFT_SLIDING_WINDOW, CAUSAL, SAME_SCHEDULE_FOR_ALL_HEADS>(params, stream); + cudaError_t status = cudaGetLastError(); + return status; +} + +template +cudaError_t BatchPrefillWithPagedKVCacheDispatched(Params& params, bool enable_pdl, + cudaStream_t stream) { + static_assert(HEAD_DIM_VO == 64 || HEAD_DIM_VO == 128 || HEAD_DIM_VO == 256); + if (MASK_MODE == MaskMode::kCustom) { + return cudaErrorNotSupported; // Not supported yet. + } + constexpr bool CAUSAL = MASK_MODE == MaskMode::kCausal; + constexpr bool MULTIITEMSCORING = MASK_MODE == MaskMode::kMultiItemScoring; + if constexpr (HEAD_DIM_QK == HEAD_DIM_VO) { + if constexpr (HEAD_DIM_VO == 64) { + // NOTE(Zihao): CTA_KV not tuned for HEAD_DIM == 64, need to optimize later + BatchPrefillWithPagedKVCacheKernelTraitsDispatched< + AttentionKernelTraits, + LEFT_SLIDING_WINDOW, CAUSAL, SAME_SCHEDULE_FOR_ALL_HEADS, Params, MULTIITEMSCORING>( + params, stream); + } else if constexpr (HEAD_DIM_VO == 128) { + BatchPrefillWithPagedKVCacheKernelTraitsDispatched< + AttentionKernelTraits, + LEFT_SLIDING_WINDOW, CAUSAL, SAME_SCHEDULE_FOR_ALL_HEADS, Params, MULTIITEMSCORING>( + params, stream); + } else { + // HEAD_DIM == 256; + // NOTE(Zihao): CTA_KV not tuned for HEAD_DIM == 256, need to optimize later + BatchPrefillWithPagedKVCacheKernelTraitsDispatched< + AttentionKernelTraits, + LEFT_SLIDING_WINDOW, CAUSAL, SAME_SCHEDULE_FOR_ALL_HEADS, Params, MULTIITEMSCORING>( + params, stream); + } + } else { + return cudaErrorNotSupported; + } + cudaError_t status = cudaGetLastError(); + return status; +}; + +} // namespace flashinfer + +#endif // FLASHINFER_ATTENTION_HOPPER_PREFILL_SM90_CUH_ diff --git a/include/flashinfer/attention/hopper/quantization/epilogue.cuh b/include/flashinfer/attention/hopper/quantization/epilogue.cuh new file mode 100644 index 0000000000000000000000000000000000000000..8bf5098d45b229b49a0d23a69338764ee271f6bd --- /dev/null +++ b/include/flashinfer/attention/hopper/quantization/epilogue.cuh @@ -0,0 +1,216 @@ +/* + * Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri + * Dao. Licensed under the BSD 3-Clause. + * + * Modified by the FlashInfer team. + */ +#ifndef FLASHINFER_ATTENTION_HOPPER_FP8_EPILOGUE_CUH_ +#define FLASHINFER_ATTENTION_HOPPER_FP8_EPILOGUE_CUH_ + +#include + +#include "../../../math.cuh" +#include "../epilogue.cuh" +#include "../named_barrier.cuh" +#include "../utils.cuh" +#include "cute/tensor.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" + +namespace flashinfer { + +using namespace cute; + +template +struct FP8CollectiveEpilogue { + using DTypeO = typename Ktraits::DTypeO; + static constexpr int CTA_Q = Ktraits::CTA_Q; + static constexpr int CTA_KV = Ktraits::CTA_KV; + static constexpr int HEAD_DIM = Ktraits::HEAD_DIM; + using TileShape_QKD = Shape, Int, Int>; + + static constexpr int NUM_WARPS = Ktraits::NUM_WARPS; + static constexpr int NUM_THREADS = NUM_WARPS * cutlass::NumThreadsPerWarp; + + static constexpr int NUM_COPY_THREADS = cutlass::NumThreadsPerWarpGroup; + static constexpr int NUM_MMA_THREADS = NUM_THREADS - NUM_COPY_THREADS; + + using SmemLayoutAtomO = decltype(cutlass::gemm::collective::detail::ss_smem_selector< + GMMA::Major::K, DTypeO, decltype(cute::get<0>(TileShape_QKD{})), + decltype(cute::get<2>(TileShape_QKD{}))>()); + using SmemLayoutO = decltype(tile_to_shape(SmemLayoutAtomO{}, select<0, 2>(TileShape_QKD{}))); + + using SmemCopyAtomO = Copy_Atom; + using SharedStorage = cute::array_aligned>; + + using ShapeT = cute::Shape; + using StrideT = cute::Shape; + using LayoutT = cute::Layout; + + using ShapeLseT = cute::Shape; + using StrideLseT = cute::Shape<_1, int64_t>; + using LayoutLseT = cute::Layout; + + using GmemTiledCopyOTMA = cute::SM90_TMA_STORE; + using TMA_O = decltype(make_tma_copy( + GmemTiledCopyOTMA{}, + make_tensor(make_gmem_ptr(static_cast(nullptr)), ShapeT{}, StrideT{}), SmemLayoutO{}, + select<0, 2>(TileShape_QKD{}), _1{})); // no mcast for O + + static constexpr int VEC_SIZE = cute::ceil_div(128, sizeof_bits_v); + static_assert(HEAD_DIM % VEC_SIZE == 0); + static constexpr int NUM_THREADS_PER_ROW = HEAD_DIM / VEC_SIZE; + static_assert(NUM_MMA_THREADS % NUM_THREADS_PER_ROW == 0); + static constexpr int NUM_ROWS = NUM_MMA_THREADS / NUM_THREADS_PER_ROW; + using TiledCopyOAtom = cute::Copy_Atom, DTypeO>; + using TiledCopyOThrLayout = decltype(cute::make_layout( + cute::make_shape(Int{}, Int{}), LayoutRight{})); + using TiledCopyOValLayout = + decltype(cute::make_layout(cute::make_shape(_1{}, Int{}), LayoutRight{})); + using TiledCopyO = + decltype(make_tiled_copy(TiledCopyOAtom{}, TiledCopyOThrLayout{}, // Thr layout + TiledCopyOValLayout{} // Val layout + )); + + // used for rmem -> smem O copy in fp8 kernel to undo column permutation + using ThreadLayoutrO = Layout, _4, _1>, Stride<_4, _32, _1, _0>>; + using ValueLayoutrO = + Layout, Int>, Stride<_0, _2, Stride<_4, _1>, _8>>; + using TiledCopyrO = decltype(make_tiled_copy(Copy_Atom, DTypeO>{}, + ThreadLayoutrO{}, ValueLayoutrO{})); + using TiledCopyShaperO = Shape<_8, Int, _16, Int>; + using SmemLayoutrO = decltype(composition(SmemLayoutO{}, Layout{})); + + // Host side kernel arguments + struct Arguments { + DTypeO* O_ptr; + LayoutT const layout_O; + float* lse_ptr; + LayoutLseT const layout_LSE; + }; + + // Device side kernel params + struct Params { + DTypeO* O_ptr; + LayoutT const layout_O; + float* lse_ptr; + LayoutLseT const layout_LSE; + }; + + static Params to_underlying_arguments(Arguments const& args) { + Tensor mO = make_tensor(make_gmem_ptr(args.O_ptr), args.layout_O); + return {args.O_ptr, args.layout_O, args.lse_ptr, args.layout_LSE}; + } + + /// Issue Tma Descriptor Prefetch -- ideally from a single thread for best performance + CUTLASS_DEVICE + static void prefetch_tma_descriptors(Params const& epilogue_params) {} + + template + CUTLASS_DEVICE void store(Params const& epilogue_params, FrgTensorO const& tOrO, + FrgTensorLSE const& lse, SharedStorage& shared_storage, + TiledMma tiled_mma, int thread_idx, BlockCoord const& block_coord) { + auto [qo_tile_idx, qo_head_idx, kv_head_idx, qo_indptr, kv_indptr, qo_len, kv_len, batch_idx] = + block_coord; + + // No need for FP8 column permutation + // as it has been done in the Transpose Phase. + Tensor sO = make_tensor(make_smem_ptr(shared_storage.smem_o.data()), SmemLayoutO{}); + auto smem_tiled_copy_O = make_tiled_copy_C(SmemCopyAtomO{}, tiled_mma); + auto smem_thr_copy_O = smem_tiled_copy_O.get_thread_slice(thread_idx); + + Tensor tOrO_out = convert_type(tOrO); + Tensor taccOrO = smem_thr_copy_O.retile_S(tOrO_out); // ((Atom,AtomNum), MMA_M, MMA_N) + Tensor taccOsO = smem_thr_copy_O.partition_D(sO); // ((Atom,AtomNum),PIPE_M,PIPE_N) + + // Make sure all WGs have finished reading V + cutlass::arch::NamedBarrier::sync(NUM_MMA_THREADS, + /*id=*/static_cast(NamedBarriers::kValueEmpty)); + cute::copy(smem_tiled_copy_O, taccOrO, taccOsO); + cutlass::arch::fence_view_async_shared(); // ensure smem writes are visible to TMA + Tensor mLSE = make_tensor(make_gmem_ptr(epilogue_params.lse_ptr), epilogue_params.layout_LSE); + Tensor gLSE = get_lse_local_tile_tensor(mLSE, Shape>{}, qo_head_idx, qo_indptr, + qo_len)(_, qo_tile_idx); + Tensor caccO = cute::make_identity_tensor(select<0, 2>(TileShape_QKD{})); + auto thread_mma = tiled_mma.get_thread_slice(thread_idx); + Tensor taccOcO = thread_mma.partition_C(caccO); // (MMA,MMA_M,MMA_K) + static_assert(decltype(size<0, 0>(taccOcO))::value == 2); + static_assert(decltype(size<0, 1>(taccOcO))::value == 2); + // taccOcO has shape ((2, 2, V), MMA_M, MMA_K), we only take only the row indices. + Tensor taccOcO_row = taccOcO(make_coord(_0{}, _, _0{}), _, _0{}); + CUTE_STATIC_ASSERT_V(size(lse) == size(taccOcO_row)); // MMA_M + if (epilogue_params.lse_ptr) { // don't write to LSE if it's nullptr + if (get<1>(taccOcO_row(_0{})) == 0) { +#pragma unroll + for (int mi = 0; mi < size(lse); ++mi) { + const int row = get<0>(taccOcO_row(mi)); + if (row < qo_len - qo_tile_idx * CTA_Q) { + gLSE(row) = lse(mi); + } + } + } + } + + // make sure all WG finish STSM o + cutlass::arch::NamedBarrier::sync(NUM_MMA_THREADS, + cutlass::arch::ReservedNamedBarriers::EpilogueBarrier); + + TiledCopyO gmem_tiled_copy_O; + int write_warp_idx = NUM_WARPS - 1; + write_O(epilogue_params.O_ptr, gmem_tiled_copy_O, epilogue_params.layout_O, + select<0, 2>(TileShape_QKD{}), sO, thread_idx, qo_tile_idx, + qo_head_idx, qo_indptr, qo_len, write_warp_idx); + } + + CUTLASS_DEVICE void store_tail() { + // tma_store_wait<0>(); + } + + // Write 0 to output and -inf to LSE + template + CUTLASS_DEVICE void store_zero(Params const& epilogue_params, SharedStorage& shared_storage, + int thread_idx, BlockCoord const& block_coord) { + auto [qo_tile_idx, qo_head_idx, kv_head_idx, qo_indptr, kv_indptr, qo_len, kv_len, batch_idx] = + block_coord; + Tensor mO = make_tensor(make_gmem_ptr(epilogue_params.O_ptr), epilogue_params.layout_O); + Tensor gO = get_local_tile_tensor(mO, select<0, 2>(TileShape_QKD{}), qo_head_idx, qo_indptr, + qo_len)(_, _, qo_tile_idx); // (O, D) + Tensor cO = cute::make_identity_tensor(gO.shape()); // (O, D) -> (o_idx, d_idx) + Tensor mLSE = make_tensor(make_gmem_ptr(epilogue_params.lse_ptr), epilogue_params.layout_LSE); + Tensor gLSE = get_lse_local_tile_tensor(mLSE, Shape>{}, qo_head_idx, qo_indptr, + qo_len)(_, qo_tile_idx); + + TiledCopyO tiled_copy_O; + auto thr_copy_O = tiled_copy_O.get_thread_slice(thread_idx); + Tensor tOgO = thr_copy_O.partition_D(gO); // (CPY, CPY_O, CPY_D) + Tensor tOrO = make_fragment_like(tOgO); // (CPY, CPY_O, CPY_D) + clear(tOrO); + Tensor tOcO = thr_copy_O.partition_D(cO); // (CPY, CPY_O, CPY_D) + Tensor tOgOGroup = flatten_1(tOgO); // (CPY, (CPY_O, CPY_D)) + Tensor tOrOGroup = flatten_1(tOrO); // (CPY, (CPY_O, CPY_D)) + Tensor tOcOGroup = flatten_1(tOcO); // (CPY, (CPY_O, CPY_D)) + + const int qo_tile_size = get<0>(TileShape_QKD{}); + int valid_qo_tile_size = std::min(qo_len - qo_tile_idx * qo_tile_size, qo_tile_size); + if (valid_qo_tile_size == qo_tile_size) { + copy(tiled_copy_O, tOrOGroup, tOgOGroup); + } else { + auto predicate_fn = [&](auto coords) { + auto s_coords = tOcOGroup(_0{}, coords); + return elem_less(get<0>(s_coords), valid_qo_tile_size); + }; + copy_if(tiled_copy_O, predicate_fn, tOrOGroup, tOgOGroup); + } + + static_assert(CTA_Q <= NUM_MMA_THREADS); + if (epilogue_params.lse_ptr) { // don't write to LSE if it's nullptr + if (thread_idx < qo_len - qo_tile_idx * CTA_Q) { + gLSE(thread_idx) = -math::inf; + } + } + } +}; + +} // namespace flashinfer + +#endif // FLASHINFER_ATTENTION_HOPPER_FP8_EPILOGUE_CUH_ diff --git a/include/flashinfer/attention/hopper/quantization/kernel_traits.cuh b/include/flashinfer/attention/hopper/quantization/kernel_traits.cuh new file mode 100644 index 0000000000000000000000000000000000000000..da5b3da9641c8624ecace2b69012784ea07ec69d --- /dev/null +++ b/include/flashinfer/attention/hopper/quantization/kernel_traits.cuh @@ -0,0 +1,246 @@ +/* + * Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri + * Dao. Licensed under the BSD 3-Clause. + * + * Modified by the FlashInfer team. + */ +#ifndef FLASHINFER_ATTENTION_HOPPER_FP8_KERNEL_TRAITS_CUH_ +#define FLASHINFER_ATTENTION_HOPPER_FP8_KERNEL_TRAITS_CUH_ + +#include + +#include "../../../cutlass_utils.cuh" +#include "cute/algorithm/copy.hpp" +#include "cute/atom/mma_atom.hpp" +#include "cutlass/cutlass.h" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/layout/layout.h" +#include "cutlass/numeric_types.h" +#include "cutlass/pipeline/pipeline.hpp" + +namespace flashinfer { + +using namespace cute; + +/* + Add additional smem for Vt + NOTE(Yilong): Should modify the mainloop to leverage the smem_v_read's early release +*/ +template +struct SharedStorageQKVOVt { + cute::array_aligned> smem_q; + cute::array_aligned> smem_k; + cute::array_aligned> smem_v; + union { + cute::array_aligned> smem_vt; + cute::array_aligned> smem_o; + }; + struct { + cutlass::arch::ClusterTransactionBarrier barrier_Q; + cutlass::arch::ClusterBarrier barrier_O; + typename MainloopPipeline::SharedStorage pipeline_k; + typename MainloopPipeline::SharedStorage pipeline_v; + // vt only use ldmatrix, which do not need TMA Pipeline + typename cutlass::PipelineAsync::SharedStorage pipeline_vt; + }; +}; + +/* + In-kernel FP8 transpose adopted from FlashAttention-3 template + https://github.com/Dao-AILab/flash-attention/blob/c7f32a8409e52a84bd8046afe7060da33036f9a5/hopper/kernel_traits.h#L217 +*/ +template +struct TranposeTraits_64x64 { + using TransposeShapeAtom_ = Shape<_64, _64>; + using TransElement = Element; + static_assert(cutlass::sizeof_bits_v == 8); + + using SmemShapeLDSM = Shape, Shape<_16, _4>>; + using SmemShapeSTSM = Shape, Shape<_16, _4>>; + + using SmemLayoutAtomV = + decltype(tile_to_shape(GMMA::Layout_K_SW64_Atom{}, TransposeShapeAtom_{})); + using SmemLayoutV = decltype(tile_to_shape( + SmemLayoutAtomV{}, + make_shape(get<1>(TileShape_QKD{}), get<2>(TileShape_QKD{}), Int{}))); + using SmemLayoutDivideV = decltype(tiled_divide(SmemLayoutV{}, TransposeShapeAtom_{})); + using FactoringShapeV = + decltype(make_shape(SmemShapeLDSM{}, shape<1>(SmemLayoutDivideV{}), + shape<2>(SmemLayoutDivideV{}), shape<3>(SmemLayoutDivideV{}))); + using SmemLayoutVTransposeSrc = + decltype(composition(SmemLayoutDivideV{}, make_layout(FactoringShapeV{}))); + + using SmemLayoutAtomVt = + decltype(tile_to_shape(GMMA::Layout_K_SW64_Atom{}, TransposeShapeAtom_{})); + // k-major Vt as target layout. this changes the memory + using SmemLayoutVt = decltype(tile_to_shape( + SmemLayoutAtomVt{}, + make_shape(get<2>(TileShape_QKD{}), get<1>(TileShape_QKD{}), Int{}))); + using SmemLayoutVtTrans = decltype(composition( + SmemLayoutVt{}, make_ordered_layout(product_each(shape(SmemLayoutV{})), Step<_2, _1, _3>{}))); + using SmemLayoutDivideVt = decltype(tiled_divide(SmemLayoutVtTrans{}, TransposeShapeAtom_{})); + using FactoringShapeVt = + decltype(make_shape(SmemShapeSTSM{}, shape<1>(SmemLayoutDivideVt{}), + shape<2>(SmemLayoutDivideVt{}), shape<3>(SmemLayoutDivideVt{}))); + using SmemLayoutVtTransposeTgt = + decltype(composition(SmemLayoutDivideVt{}, make_layout(FactoringShapeVt{}))); +}; + +/* + In-kernel Transpose of smemV into smemVt with ldmatrix.trans & stmatrix. + Note that all magic number corresponds to the /quantization/kernel_traits.cuh setup. + This transpose is not a general transpose, but a specific one for the FP8 MMA_PV: + 1. K-dimension: (2,2,4,4):(1,8,2,16), which adheres to the accum_P's layout + 2. N-dimension: (8,2,4):(2,1,16), which needs repermutation when rmemO -> smemO +*/ +template +struct SmemTransposeFP8_64x64 { + using Element = typename Ktraits::DTypeKV; + using SmemLayoutVTransposeSrc = typename Ktraits::SmemLayoutVTransposeSrc; + using SmemLayoutVtTransposeTgt = typename Ktraits::SmemLayoutVtTransposeTgt; + static_assert(cutlass::sizeof_bits_v == 8); + + using ldsm_thread_shape = Shape<_4, _1, _8, _4>; + using ldsm_value_shape = Shape<_2, _8, _2, _1>; + using ldsm_value_stride = Stride<_2, _4, _1, _0>; + // use trans to do 16bits transpose + // which needs permutation to separate 8bits row and column + using TiledCopyLDSM = + decltype(make_tiled_copy(Copy_Atom{}, Layout{}, + Layout{})); + TiledCopyLDSM tiled_copy_ldsm; + + using stsm_thread_shape = Shape<_4, _1, _8, _4>; + using stsm_value_shape = Shape<_4, _4, _2, _1>; + using stsm_value_stride = Stride<_1, _8, _4, _0>; + + using TiledCopySTSM = + decltype(make_tiled_copy(Copy_Atom{}, Layout{}, + Layout{})); + TiledCopySTSM tiled_copy_stsm; + + template + CUTLASS_DEVICE void _tranpose(SmemTensor&& s_in, SmemTensorOut&& s_out) { + using namespace cute; + + auto tid = threadIdx.x; + auto thr_copy_ldsm = tiled_copy_ldsm.get_thread_slice(tid); + auto thr_copy_stsm = tiled_copy_stsm.get_thread_slice(tid); + + auto tXsX = thr_copy_ldsm.partition_S(s_in); + auto tXrX = make_tensor(shape(tXsX)); + auto tXsX_out = thr_copy_stsm.partition_D(s_out); + + cute::copy(tiled_copy_ldsm, tXsX, tXrX); + auto data = tXrX.data(); + CUTLASS_PRAGMA_UNROLL + for (int n = 0; n < size(tXrX); n += 8) { + uint32_t* data_32bit = reinterpret_cast(&data[n]); + auto upper = data_32bit[0]; + auto lower = data_32bit[1]; + // select row-major elements. + // from (0 1 16 17) (128 129 144 145) to (0 16 128 144) (1 17 129 145) + // which is (0 1 8 9) + data_32bit[0] = __byte_perm(upper, lower, 0x6420); + data_32bit[1] = __byte_perm(upper, lower, 0x7531); + } + cute::copy(tiled_copy_stsm, tXrX, tXsX_out); + } + + template + CUTLASS_DEVICE void do_transpose(SmemTensor& s_in, SmemTensorOut& s_out, int stage_idx) { + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < shape<2>(SmemLayoutVTransposeSrc{}); ++j) { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < shape<1>(SmemLayoutVTransposeSrc{}); ++i) { + this->_tranpose(flatten(s_in(_, i, j, stage_idx)), flatten(s_out(_, i, j, stage_idx))); + } + } + // For FP8 kernel, all WG threads will arrive for issuing ldmatrix + cutlass::arch::NamedBarrier::sync(Ktraits::NUM_PRODUCER_THREADS, + static_cast(NamedBarriers::kProducerWG) /*id*/); + } +}; + +template +struct FP8AttentionKernelTraits { + using AttentionVariant = AttentionVariant_; + + using DTypeQ = DTypeQ_; + using DTypeKV = DTypeKV_; + using DTypeO = DTypeO_; + using IdType = IdType_; + using DTypeQKAccum = float; + + static constexpr int CTA_Q = CTA_Q_; + static_assert(CTA_Q % 64 == 0); + static constexpr int CTA_KV = CTA_KV_; + static constexpr int HEAD_DIM = HEAD_DIM_; + static_assert(HEAD_DIM % 32 == 0); + + static constexpr int NUM_WARPS = ((CTA_Q / 64) + 1) * 4; + static constexpr int NUM_THREADS = NUM_WARPS * cutlass::NumThreadsPerWarp; + // NOTE(Zihao): the following constant should only be used when TMA is enabled, + // where only one warp inside a warp group is used for TMA. + + // In FP16 kernel, only one thread of single warp within the producer WG is working + // For FP8, we use the entire WG for tranposing V + static constexpr int NUM_PRODUCER_THREADS = cutlass::NumThreadsPerWarpGroup; + + using TileShape_QKD = Shape, Int, Int>; + + static constexpr int NUM_STAGES = NUM_STAGES_; + + using AtomLayoutQKD = Layout, _1, _1>>; + using TiledMmaQK = decltype(cute::make_tiled_mma( + cute::GMMA::ss_op_selector(), AtomLayoutQKD{})); + + // FP8 needs K-major for both P / V + using TiledMmaPV = decltype(cute::make_tiled_mma( + cute::GMMA::rs_op_selector(TileShape_QKD{})), GMMA::Major::K, + GMMA::Major::K>(), + AtomLayoutQKD{})); + + static constexpr int NUM_MMA_THREADS = size(TiledMmaQK{}); + + using SmemLayoutAtomQ = decltype(cutlass::gemm::collective::detail::ss_smem_selector< + GMMA::Major::K, DTypeQ, decltype(cute::get<0>(TileShape_QKD{})), + decltype(cute::get<2>(TileShape_QKD{}))>()); + using SmemLayoutQ = decltype(tile_to_shape(SmemLayoutAtomQ{}, select<0, 2>(TileShape_QKD{}))); + + using SmemLayoutAtomK = decltype(cutlass::gemm::collective::detail::ss_smem_selector< + GMMA::Major::K, DTypeKV, decltype(cute::get<1>(TileShape_QKD{})), + decltype(cute::get<2>(TileShape_QKD{}))>()); + using SmemLayoutK = decltype(tile_to_shape( + SmemLayoutAtomK{}, + make_shape(shape<1>(TileShape_QKD{}), shape<2>(TileShape_QKD{}), Int{}))); + + using VTranposeTraits = TranposeTraits_64x64; + using SmemLayoutV = typename VTranposeTraits::SmemLayoutV; + using SmemLayoutVt = typename VTranposeTraits::SmemLayoutVt; + using SmemLayoutVTransposeSrc = typename VTranposeTraits::SmemLayoutVTransposeSrc; + using SmemLayoutVtTransposeTgt = typename VTranposeTraits::SmemLayoutVtTransposeTgt; + + using SmemLayoutAtomO = decltype(cutlass::gemm::collective::detail::ss_smem_selector< + GMMA::Major::K, DTypeO, decltype(cute::get<0>(TileShape_QKD{})), + decltype(cute::get<2>(TileShape_QKD{}))>()); + using SmemLayoutO = decltype(tile_to_shape(SmemLayoutAtomO{}, select<0, 2>(TileShape_QKD{}))); + using MainloopPipeline = + std::conditional_t, + typename cutlass::PipelineAsync>; + using MainloopPipelineNoTMA = typename cutlass::PipelineAsync; + using PipelineState = typename cutlass::PipelineState; + + // Modify SharedStorage + using SharedStorage = + SharedStorageQKVOVt; +}; + +} // namespace flashinfer + +#endif // FLASHINFER_ATTENTION_HOPPER_FP8_KERNEL_TRAITS_CUH_ diff --git a/include/flashinfer/attention/hopper/quantization/mainloop_load.cuh b/include/flashinfer/attention/hopper/quantization/mainloop_load.cuh new file mode 100644 index 0000000000000000000000000000000000000000..988f7e9acac9efbf36be47be6d17664a94bf5d09 --- /dev/null +++ b/include/flashinfer/attention/hopper/quantization/mainloop_load.cuh @@ -0,0 +1,321 @@ +/* + * Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri + * Dao. Licensed under the BSD 3-Clause. + * + * Modified by the FlashInfer team. + */ +#ifndef FLASHINFER_ATTENTION_HOPPER_FP8_MAINLOOP_CUH_ +#define FLASHINFER_ATTENTION_HOPPER_FP8_MAINLOOP_CUH_ + +#include +#include +#include +#include + +#include "../../../math.cuh" +#include "../named_barrier.cuh" +#include "../utils.cuh" +#include "cute/tensor.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/pipeline/pipeline.hpp" +#include "kernel_traits.cuh" + +namespace flashinfer { + +using namespace cute; + +template +struct FP8CollectiveMainloop { + using DTypeQ = typename Ktraits::DTypeQ; + using DTypeKV = typename Ktraits::DTypeKV; + using TileShape_QKD = typename Ktraits::TileShape_QKD; + static constexpr int CTA_Q = get<0>(TileShape_QKD{}); + static constexpr int CTA_KV = get<1>(TileShape_QKD{}); + + static constexpr int NUM_STAGES = Ktraits::NUM_STAGES; + static constexpr int NUM_MMA_THREADS = Ktraits::NUM_MMA_THREADS; + static constexpr int HEAD_DIM = Ktraits::HEAD_DIM; + + using GmemTiledCopyQ = cute::SM90_TMA_LOAD; + using GmemTiledCopyKV = cute::SM90_TMA_LOAD; + + using SmemLayoutQ = typename Ktraits::SmemLayoutQ; + using SmemLayoutK = typename Ktraits::SmemLayoutK; + using SmemLayoutV = typename Ktraits::SmemLayoutV; + using SmemLayoutVt = typename Ktraits::SmemLayoutVt; + + using ShapeT = cute::Shape; + using StrideT = cute::Shape; // (N, D, H) + using LayoutT = cute::Layout; + + using ShapeLseT = cute::Shape; + using StrideLseT = cute::Shape<_1, int64_t>; + using LayoutLseT = cute::Layout; + + using TMA_Q = decltype(make_tma_copy( + GmemTiledCopyQ{}, + make_tensor(make_gmem_ptr(static_cast(nullptr)), + repeat_like(StrideT{}, int32_t(0)), StrideT{}), + SmemLayoutQ{}, select<0, 2>(TileShape_QKD{}), _1{})); // no mcast for Q + + using TMA_K = decltype(make_tma_copy( + GmemTiledCopyKV{}, + make_tensor(make_gmem_ptr(static_cast(nullptr)), + repeat_like(StrideT{}, int32_t(0)), StrideT{}), + take<0, 2>(SmemLayoutK{}), select<1, 2>(TileShape_QKD{}), _1{})); // no mcast + + using TMA_V = decltype(make_tma_copy( + GmemTiledCopyKV{}, + make_tensor(make_gmem_ptr(static_cast(nullptr)), + repeat_like(StrideT{}, int32_t(0)), StrideT{}), + take<0, 2>(SmemLayoutV{}), select<1, 2>(TileShape_QKD{}), _1{})); // no mcast + + static constexpr bool USE_TMA_LOAD_KV = true; + using MainloopPipeline = typename Ktraits::MainloopPipeline; + using PipelineParams = typename MainloopPipeline::Params; + using PipelineState = typename MainloopPipeline::PipelineState; + using MainloopPipelineVt = typename Ktraits::MainloopPipelineNoTMA; + using PipelineParamsVt = typename MainloopPipelineVt::Params; + + // Set the bytes transferred in this TMA transaction (may involve multiple issues) + static constexpr uint32_t TmaTransactionBytesQ = + static_cast(size(SmemLayoutQ{}) * cutlass::sizeof_bits_v / 8); + static constexpr uint32_t TmaTransactionBytesK = + static_cast(size(take<0, 2>(SmemLayoutK{})) * cutlass::sizeof_bits_v / 8); + + // Whether use scheduler barrier or hardware warp scheduler, using heuristic based on data type + // and head dim + static constexpr bool UseSchedulerBarrier = + cutlass::sizeof_bits_v == 8 ? HEAD_DIM >= 128 : HEAD_DIM <= 128; + using WarpScheduler = WarpScheduler; + + // Host side kernel arguments + struct Arguments { + DTypeQ const* Q_ptr; + LayoutT layout_Q; + DTypeKV const* K_ptr; + LayoutT layout_K; + DTypeKV const* V_ptr; + LayoutT layout_V; + int window_left; + AdditionalParams additional_params; + }; + + // Device side kernel params + struct Params { + LayoutT layout_Q; + LayoutT layout_K; + LayoutT layout_V; + TMA_Q tma_load_Q; + TMA_K tma_load_K; + TMA_V tma_load_V; + int window_left; + AdditionalParams additional_params; + using DTypeKV = typename Ktraits::DTypeKV; + }; + + static Params to_underlying_arguments(Arguments const& args) { + Tensor mQ = make_tensor(make_gmem_ptr(args.Q_ptr), args.layout_Q); + TMA_Q tma_load_Q = make_tma_copy(GmemTiledCopyQ{}, mQ, SmemLayoutQ{}, + select<0, 2>(TileShape_QKD{}), _1{}); // no mcast for Q + Tensor mK = make_tensor(make_gmem_ptr(args.K_ptr), args.layout_K); + TMA_K tma_load_K = make_tma_copy(GmemTiledCopyKV{}, mK, SmemLayoutK{}(_, _, _0{}), + select<1, 2>(TileShape_QKD{}), _1{}); // no mcast + Tensor mV = make_tensor(make_gmem_ptr(args.V_ptr), args.layout_V); + TMA_V tma_load_V = make_tma_copy(GmemTiledCopyKV{}, mV, SmemLayoutV{}(_, _, _0{}), + select<1, 2>(TileShape_QKD{}), _1{}); // no mcast + return {args.layout_Q, args.layout_K, args.layout_V, tma_load_Q, + tma_load_K, tma_load_V, args.window_left, args.additional_params}; + } + + /// Issue Tma Descriptor Prefetch -- ideally from a single thread for best performance + CUTLASS_DEVICE + static void prefetch_tma_descriptors(Params const& mainloop_params) { + cute::prefetch_tma_descriptor(mainloop_params.tma_load_Q.get_tma_descriptor()); + cute::prefetch_tma_descriptor(mainloop_params.tma_load_K.get_tma_descriptor()); + cute::prefetch_tma_descriptor(mainloop_params.tma_load_V.get_tma_descriptor()); + } + + CUTLASS_DEVICE + int get_num_kv_tiles(Params const& mainloop_params, int q_tile_idx, const int qo_len, + const int kv_len) { + static constexpr int CTA_Q = get<0>(TileShape_QKD{}); + static constexpr int CTA_KV = get<1>(TileShape_QKD{}); + int num_kv_tiles = cute::ceil_div(kv_len, CTA_KV); + if constexpr (CAUSAL) { + num_kv_tiles = std::min(num_kv_tiles, + cute::ceil_div((q_tile_idx + 1) * CTA_Q + kv_len - qo_len, CTA_KV)); + } + + return num_kv_tiles; + } + + template + CUTLASS_DEVICE void load(Params const& mainloop_params, MainloopPipeline pipeline_k, + MainloopPipeline pipeline_v, MainloopPipelineVt pipeline_vt, + PipelineState& smem_pipe_write, PipelineState& smem_pipe_read, + SharedStorage& shared_storage, Scheduler& scheduler, + typename Scheduler::Params const& scheduler_params, + typename Scheduler::WorkTileInfo& work_tile_info, + BlockCoord const& block_coord, int work_idx) { + Tensor sQ = make_tensor(make_smem_ptr(shared_storage.smem_q.data()), SmemLayoutQ{}); + Tensor sK = make_tensor(make_smem_ptr(shared_storage.smem_k.data()), SmemLayoutK{}); + Tensor sV = make_tensor(make_smem_ptr(shared_storage.smem_v.data()), SmemLayoutV{}); + + Tensor mQ = mainloop_params.tma_load_Q.get_tma_tensor(mainloop_params.layout_Q.shape()); + Tensor mK = mainloop_params.tma_load_K.get_tma_tensor(mainloop_params.layout_K.shape()); + Tensor mV = mainloop_params.tma_load_V.get_tma_tensor(mainloop_params.layout_V.shape()); + + // *** Prepare In-kernel V Transpose *** + using SmemLayoutVTransposeSrc = typename Ktraits::SmemLayoutVTransposeSrc; + using SmemLayoutVtTransposeTgt = typename Ktraits::SmemLayoutVtTransposeTgt; + + Tensor sV_src = as_position_independent_swizzle_tensor( + make_tensor(make_smem_ptr(shared_storage.smem_v.data()), SmemLayoutVTransposeSrc{})); + Tensor sVt_tgt = as_position_independent_swizzle_tensor( + make_tensor(make_smem_ptr(shared_storage.smem_vt.data()), SmemLayoutVtTransposeTgt{})); + auto v_tranposer = SmemTransposeFP8_64x64(); + + auto [q_tile_idx, qo_head_idx, kv_head_idx, qo_indptr, kv_indptr, qo_len, kv_len, batch_idx] = + block_coord; + + // Prepare the TMA loads + Tensor gQ = get_local_tile_tensor(mQ, select<0, 2>(TileShape_QKD{}), qo_head_idx, qo_indptr, + qo_len)(_, _, q_tile_idx); // (Q, D) + Tensor gK = get_local_tile_tensor(mK, select<1, 2>(TileShape_QKD{}), kv_head_idx, kv_indptr, + kv_len); // (K, D, _) + Tensor gV = get_local_tile_tensor(mV, select<1, 2>(TileShape_QKD{}), kv_head_idx, kv_indptr, + kv_len); // (K, D, _) + + Tensor sQ_x = make_tensor(sQ.data(), make_layout(sQ.layout(), Layout<_1>{})); + Tensor gQ_x = make_tensor(gQ.data(), make_layout(gQ.layout(), Layout<_1>{})); + auto [tQgQ, tQsQ] = + tma_partition(mainloop_params.tma_load_Q, _0{}, Layout<_1>{}, group_modes<0, 2>(sQ_x), + group_modes<0, 2>(gQ_x)); // (TMA), (TMA) + auto [tKgK, tKsK] = + tma_partition(mainloop_params.tma_load_K, _0{}, Layout<_1>{}, group_modes<0, 2>(sK), + group_modes<0, 2>(gK)); // (TMA, k), (TMA, PIPE) + auto [tVgV, tVsV] = + tma_partition(mainloop_params.tma_load_V, _0{}, Layout<_1>{}, group_modes<0, 2>(sV), + group_modes<0, 2>(gV)); // (TMA, k), (TMA, PIPE) + + int num_kv_tiles = get_num_kv_tiles(mainloop_params, q_tile_idx, qo_len, kv_len); + int kv_tile_idx = num_kv_tiles - 1; + int swa_begin_kv_tile_idx = 0; + if constexpr (LEFT_SLIDING_WINDOW) { + swa_begin_kv_tile_idx = get_swa_begin_kv_tile_idx(mainloop_params.window_left, + q_tile_idx, qo_len, kv_len); + } + + // All WG proceeds here, only one thread in each WG will issue TMA load + int lane_predicate = cute::elect_one_sync(); + int warp_idx_in_warpgroup = __shfl_sync(0xffffffff, (threadIdx.x / 32) % 4, 0); + bool issue_tma_thread = (warp_idx_in_warpgroup == 0) && (lane_predicate == 1); + + if (issue_tma_thread) { + pipeline_k.producer_acquire(smem_pipe_write); + copy(mainloop_params.tma_load_K.with(*pipeline_k.producer_get_barrier(smem_pipe_write), + /*mcast_mask=*/0), + tKgK(_, kv_tile_idx), tKsK(_, smem_pipe_write.index())); + } + + // Wait for the MMA warpgroups to say that smem_q is ready + cutlass::arch::NamedBarrier::sync(NUM_MMA_THREADS + Ktraits::NUM_PRODUCER_THREADS, + static_cast(NamedBarriers::kQueryEmpty)); + + if (issue_tma_thread) { + shared_storage.barrier_Q.arrive_and_expect_tx(TmaTransactionBytesQ); + copy(mainloop_params.tma_load_Q.with( + reinterpret_cast( + shared_storage.barrier_Q), + /*mcast_mask=*/0), + tQgQ, tQsQ); + + pipeline_v.producer_acquire(smem_pipe_write); + copy(mainloop_params.tma_load_V.with(*pipeline_v.producer_get_barrier(smem_pipe_write), + /*mcast_mask=*/0), + tVgV(_, kv_tile_idx), tVsV(_, smem_pipe_write.index())); + } + + // Wait for warp 1 to signal that smem_v are ready and V can be copied from gmem + // Need ClusterBarrier, not just NamedBarrier. Otherwise we might have CTA 0 finishing the + // TMA store on O first, call TMA multicast load on V, before CTA 1 can finishing TMA store on + // O. + shared_storage.barrier_O.wait((work_idx + 1) % 2); + + pipeline_v.consumer_wait(smem_pipe_read); + pipeline_vt.producer_acquire(smem_pipe_write); + v_tranposer.do_transpose(sV_src, sVt_tgt, smem_pipe_read.index()); + pipeline_vt.producer_commit(smem_pipe_write); + pipeline_v.consumer_release(smem_pipe_read); + ++smem_pipe_read; + ++smem_pipe_write; + --kv_tile_idx; + + constexpr int num_left_iter = Ktraits::NUM_STAGES - 1; +#pragma unroll 2 + for (int iter = 0; iter < num_left_iter && kv_tile_idx >= swa_begin_kv_tile_idx; + --kv_tile_idx, ++iter) { + if (issue_tma_thread) { + pipeline_k.producer_acquire(smem_pipe_write); + copy(mainloop_params.tma_load_K.with(*pipeline_k.producer_get_barrier(smem_pipe_write), + /*mcast_mask=*/0), + tKgK(_, kv_tile_idx), tKsK(_, smem_pipe_write.index())); + + pipeline_v.producer_acquire(smem_pipe_write); + copy(mainloop_params.tma_load_V.with(*pipeline_v.producer_get_barrier(smem_pipe_write), + /*mcast_mask=*/0), + tVgV(_, kv_tile_idx), tVsV(_, smem_pipe_write.index())); + } + + pipeline_v.consumer_wait(smem_pipe_read); + pipeline_vt.producer_acquire(smem_pipe_write); + v_tranposer.do_transpose(sV_src, sVt_tgt, smem_pipe_read.index()); + pipeline_vt.producer_commit(smem_pipe_write); + pipeline_v.consumer_release(smem_pipe_read); + ++smem_pipe_read; + ++smem_pipe_write; + } + +#pragma unroll 2 + for (; kv_tile_idx >= swa_begin_kv_tile_idx; --kv_tile_idx) { + if (issue_tma_thread) { + pipeline_k.producer_acquire(smem_pipe_write); + copy(mainloop_params.tma_load_K.with(*pipeline_k.producer_get_barrier(smem_pipe_write), + /*mcast_mask=*/0), + tKgK(_, kv_tile_idx), tKsK(_, smem_pipe_write.index())); + + pipeline_v.producer_acquire(smem_pipe_write); + copy(mainloop_params.tma_load_V.with(*pipeline_v.producer_get_barrier(smem_pipe_write), + /*mcast_mask=*/0), + tVgV(_, kv_tile_idx), tVsV(_, smem_pipe_write.index())); + } + pipeline_v.consumer_wait(smem_pipe_read); + pipeline_vt.producer_acquire(smem_pipe_write); + v_tranposer.do_transpose(sV_src, sVt_tgt, smem_pipe_read.index()); + pipeline_vt.producer_commit(smem_pipe_write); + pipeline_v.consumer_release(smem_pipe_read); + ++smem_pipe_read; + ++smem_pipe_write; + } + scheduler.prefetch_next_work(scheduler_params, work_tile_info); + scheduler.broadcast_next_work(work_tile_info); + } + + CUTLASS_DEVICE void load_tail(MainloopPipeline pipeline_k, MainloopPipeline pipeline_v, + PipelineState& smem_pipe_write) { + // This func is not useful as blocking transpose is enabled + // WG will not early exit + int lane_predicate = cute::elect_one_sync(); + int warp_idx_in_warpgroup = __shfl_sync(0xffffffff, (threadIdx.x / 32) % 4, 0); + if (warp_idx_in_warpgroup == 0 && lane_predicate) { + pipeline_k.producer_tail(smem_pipe_write); + pipeline_v.producer_tail(smem_pipe_write); + } + } +}; + +} // namespace flashinfer + +#endif // FLASHINFER_ATTENTION_HOPPER_FP8_MAINLOOP_CUH_ diff --git a/include/flashinfer/attention/hopper/quantization/mainloop_mma.cuh b/include/flashinfer/attention/hopper/quantization/mainloop_mma.cuh new file mode 100644 index 0000000000000000000000000000000000000000..9720af575edb5e8f911f4e981c7adbc82ed44e64 --- /dev/null +++ b/include/flashinfer/attention/hopper/quantization/mainloop_mma.cuh @@ -0,0 +1,240 @@ +/* + * Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri + * Dao. Licensed under the BSD 3-Clause. + * + * Modified by the FlashInfer team. + */ +#ifndef FLASHINFER_ATTENTION_HOPPER_FP8_MAINLOOP_MMA_CUH_ +#define FLASHINFER_ATTENTION_HOPPER_FP8_MAINLOOP_MMA_CUH_ + +#include +#include +#include +#include + +namespace flashinfer { + +template +CUTLASS_DEVICE void mma_fp8(const Params& mainloop_params, AttentionVariant& variant, + MainloopPipeline pipeline_k, MainloopPipelineVt pipeline_vt, + PipelineState& smem_pipe_read_k, PipelineState& smem_pipe_read_v, + FrgTensorO& tOrO, AttentionUpdater& attention_updater, + int kv_tile_idx_count, int swa_begin_kv_tile_idx, + int swa_end_kv_tile_idx, int thread_idx, int work_idx, int q_tile_idx, + SharedStorage& shared_storage, const int32_t qo_len, + const int32_t kv_len, const int32_t qo_head_idx, + const int32_t kv_head_idx, const int32_t batch_idx) { + using DTypeQ = typename Ktraits::DTypeQ; + using DTypeKV = typename Ktraits::DTypeKV; + using IdType = typename Ktraits::IdType; + using TileShape_QKD = typename Ktraits::TileShape_QKD; + static constexpr int NUM_MMA_THREADS = Ktraits::NUM_MMA_THREADS; + using SmemLayoutQ = typename Ktraits::SmemLayoutQ; + using SmemLayoutK = typename Ktraits::SmemLayoutK; + using SmemLayoutV = typename Ktraits::SmemLayoutV; + using SmemLayoutVt = typename Ktraits::SmemLayoutVt; + static_assert(is_rmem::value, "O tensor must be rmem resident."); + + static constexpr int CTA_Q = get<0>(TileShape_QKD{}); + static constexpr int CTA_KV = get<1>(TileShape_QKD{}); + + Tensor sQ = make_tensor(make_smem_ptr(shared_storage.smem_q.data()), SmemLayoutQ{}); + Tensor sK = make_tensor(make_smem_ptr(shared_storage.smem_k.data()), SmemLayoutK{}); + Tensor sVt = make_tensor(make_smem_ptr(shared_storage.smem_vt.data()), SmemLayoutVt{}); + + typename Ktraits::TiledMmaQK tiled_mma_qk; + typename Ktraits::TiledMmaPV tiled_mma_pv; + auto threadMmaQK = tiled_mma_qk.get_thread_slice(thread_idx); + auto threadMmaPV = tiled_mma_pv.get_thread_slice(thread_idx); + + Tensor tSrQ = threadMmaQK.partition_fragment_A(sQ); + Tensor tSrK = threadMmaQK.partition_fragment_B(sK); + Tensor tOrV = threadMmaPV.partition_fragment_B(sVt); + + auto consumer_wait = [](auto& pipeline, auto& smem_pipe_read) { + auto barrier_token = pipeline.consumer_try_wait(smem_pipe_read); + pipeline.consumer_wait(smem_pipe_read, barrier_token); + }; + + tiled_mma_pv.accumulate_ = GMMA::ScaleOut::Zero; + int kv_tile_idx = kv_tile_idx_count - 1; + + cutlass::ConsumerToken barrier_token = + static_cast(shared_storage.barrier_Q.try_wait(work_idx % 2)); + if (barrier_token == cutlass::BarrierStatus::WaitAgain) { + shared_storage.barrier_Q.wait(work_idx % 2); + } + + Tensor tSrS = partition_fragment_C(tiled_mma_qk, select<0, 1>(TileShape_QKD{})); + consumer_wait(pipeline_k, smem_pipe_read_k); + + WarpScheduler::barrier_sync(); + gemm(tiled_mma_qk, tSrQ, tSrK(_, _, _, smem_pipe_read_k.index()), + tSrS); + WarpScheduler::barrier_arrive(); + + if (work_idx != 0) { + int lane_predicate = cute::elect_one_sync(); + if (cutlass::canonical_warp_idx_sync() == Ktraits::NUM_WARPS - 1 && lane_predicate) { +#pragma unroll + for (uint32_t cta_id = 0; cta_id < 1; ++cta_id) { + shared_storage.barrier_O.arrive(cta_id, lane_predicate); + } + } + } + warpgroup_wait<0>(); + pipeline_k.consumer_release(smem_pipe_read_k); + ++smem_pipe_read_k; + + auto col_limit_right = [&](int qo_idx) { return qo_idx + 1 + kv_len - qo_len; }; + auto col_limit_left = [&](int qo_idx) { + return qo_idx + kv_len - qo_len - mainloop_params.window_left; + }; + { + Tensor cS = cute::make_identity_tensor(select<0, 1>(TileShape_QKD{})); + Tensor tScS = threadMmaQK.partition_C(cS); +#pragma unroll + for (int i = 0; i < size(tSrS); ++i) { + int qo_idx = get<0>(tScS(i)) + q_tile_idx * CTA_Q; + int kv_idx = get<1>(tScS(i)) + kv_tile_idx * CTA_KV; + tSrS(i) = variant.LogitsTransform(mainloop_params, tSrS(i), /*batch_idx=*/batch_idx, qo_idx, + kv_idx, qo_head_idx, kv_head_idx); + if constexpr (!CAUSAL) { // Just masking based on col + if (kv_idx >= kv_len) { + tSrS(i) = AttentionUpdater::fill_value; + } + } else { + if (kv_idx >= std::min(kv_len, col_limit_right(qo_idx))) { + tSrS(i) = AttentionUpdater::fill_value; + } + } + if constexpr (LEFT_SLIDING_WINDOW) { + if (kv_idx < col_limit_left(qo_idx)) { + tSrS(i) = AttentionUpdater::fill_value; + } + } + } + } + + attention_updater.update(tSrS); + // Re-quantize P after softmax + variant.PQuantize(tSrS); + + // Cast back to FP8 + Tensor tOrP = + make_tensor(convert_type(tSrS).data(), convert_layout_acc_Aregs_fp8(tSrS.layout())); + permute_regs_A_to_C(tOrP); + + constexpr int n_masking_steps = CAUSAL ? cute::ceil_div(CTA_Q, CTA_KV) : 0; + // masking loops +#pragma unroll + for (int masking_step = 0; masking_step < n_masking_steps && kv_tile_idx > swa_begin_kv_tile_idx; + ++masking_step, --kv_tile_idx) { + Tensor tSrS = partition_fragment_C(tiled_mma_qk, select<0, 1>(TileShape_QKD{})); + consumer_wait(pipeline_k, smem_pipe_read_k); + WarpScheduler::barrier_sync(); + gemm(tiled_mma_qk, tSrQ, tSrK(_, _, _, smem_pipe_read_k.index()), + tSrS); + if (masking_step > 0) { + attention_updater.rescale_o(tOrO); + } + consumer_wait(pipeline_vt, smem_pipe_read_v); + gemm(tiled_mma_pv, tOrP, + tOrV(_, _, _, smem_pipe_read_v.index()), tOrO); + WarpScheduler::barrier_arrive(); + warpgroup_wait<1>(); + pipeline_k.consumer_release(smem_pipe_read_k); // release K + Tensor cS = cute::make_identity_tensor(select<0, 1>(TileShape_QKD{})); + Tensor tScS = threadMmaQK.partition_C(cS); +#pragma unroll + for (int i = 0; i < size(tSrS); ++i) { + int qo_idx = get<0>(tScS(i)) + q_tile_idx * CTA_Q; + int kv_idx = get<1>(tScS(i)) + (kv_tile_idx - 1) * CTA_KV; + tSrS(i) = variant.LogitsTransform(mainloop_params, tSrS(i), /*batch_idx=*/batch_idx, qo_idx, + kv_idx, qo_head_idx, kv_head_idx); + if (kv_idx >= col_limit_right(qo_idx)) { + tSrS(i) = AttentionUpdater::fill_value; + } + if constexpr (LEFT_SLIDING_WINDOW) { + if (kv_idx < col_limit_left(qo_idx)) { + tSrS(i) = AttentionUpdater::fill_value; + } + } + } + attention_updater.update(tSrS); + // Re-quantize P after softmax + variant.PQuantize(tSrS); + + warpgroup_wait<0>(); + pipeline_vt.consumer_release(smem_pipe_read_v); // release V + ++smem_pipe_read_k; + ++smem_pipe_read_v; + cute::copy(make_tensor(convert_type(tSrS).data(), + convert_layout_acc_Aregs_fp8(tSrS.layout())), + tOrP); + permute_regs_A_to_C(tOrP); + } + +#pragma unroll 1 + for (; kv_tile_idx > swa_end_kv_tile_idx + 1; --kv_tile_idx) { + Tensor tSrS = partition_fragment_C(tiled_mma_qk, select<0, 1>(TileShape_QKD{})); + consumer_wait(pipeline_k, smem_pipe_read_k); + WarpScheduler::barrier_sync(); + gemm(tiled_mma_qk, tSrQ, tSrK(_, _, _, smem_pipe_read_k.index()), + tSrS); + attention_updater.rescale_o(tOrO); + consumer_wait(pipeline_vt, smem_pipe_read_v); + gemm(tiled_mma_pv, tOrP, + tOrV(_, _, _, smem_pipe_read_v.index()), tOrO); + WarpScheduler::barrier_arrive(); + warpgroup_wait<1>(); + pipeline_k.consumer_release(smem_pipe_read_k); // release K + // #pragma unroll + Tensor cS = cute::make_identity_tensor(select<0, 1>(TileShape_QKD{})); + Tensor tScS = threadMmaQK.partition_C(cS); +#pragma unroll + for (int i = 0; i < size(tSrS); ++i) { + int qo_idx = get<0>(tScS(i)) + q_tile_idx * CTA_Q; + int kv_idx = get<1>(tScS(i)) + (kv_tile_idx - 1) * CTA_KV; + tSrS(i) = variant.LogitsTransform(mainloop_params, tSrS(i), /*batch_idx=*/batch_idx, qo_idx, + kv_idx, qo_head_idx, kv_head_idx); + } + + attention_updater.update(tSrS); + // Re-quantize P after softmax + variant.PQuantize(tSrS); + + warpgroup_wait<0>(); + pipeline_vt.consumer_release(smem_pipe_read_v); // release V + ++smem_pipe_read_k; + ++smem_pipe_read_v; + cute::copy(make_tensor(convert_type(tSrS).data(), + convert_layout_acc_Aregs_fp8(tSrS.layout())), + tOrP); + permute_regs_A_to_C(tOrP); + } + + // Tell warp 0 that smem_q is ready + cutlass::arch::NamedBarrier::arrive(NUM_MMA_THREADS + Ktraits::NUM_PRODUCER_THREADS, + /*id=*/static_cast(NamedBarriers::kQueryEmpty)); + attention_updater.rescale_o(tOrO); + consumer_wait(pipeline_vt, smem_pipe_read_v); + gemm(tiled_mma_pv, tOrP, tOrV(_, _, _, smem_pipe_read_v.index()), + tOrO); + attention_updater.finalize(tSrS, variant.scale_pv); + warpgroup_wait<0>(); + pipeline_vt.consumer_release(smem_pipe_read_v); // release V, otherwise producers will hang + ++smem_pipe_read_v; + + attention_updater.rescale_o(tOrO); + // Dequantize output o with P/V scale + variant.ODequantize(mainloop_params, tOrO, qo_head_idx, kv_head_idx); + return; +} + +} // namespace flashinfer + +#endif // FLASHINFER_ATTENTION_HOPPER_FP8_MAINLOOP_MMA_CUH_ diff --git a/include/flashinfer/attention/hopper/quantization/mainloop_sparse_load.cuh b/include/flashinfer/attention/hopper/quantization/mainloop_sparse_load.cuh new file mode 100644 index 0000000000000000000000000000000000000000..1b17842443868710cd14aeaba79066d6c3e25010 --- /dev/null +++ b/include/flashinfer/attention/hopper/quantization/mainloop_sparse_load.cuh @@ -0,0 +1,366 @@ +/* + * Copyright (c) 2024 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef FLASHINFER_ATTENTION_HOPPER_FP8_SPARSE_MAINLOOP_CUH_ +#define FLASHINFER_ATTENTION_HOPPER_FP8_SPARSE_MAINLOOP_CUH_ + +#include +#include +#include +#include + +#include +#include +#include + +#include "../../../math.cuh" +#include "../block_sparse_gather.cuh" +#include "../named_barrier.cuh" +#include "../utils.cuh" +#include "kernel_traits.cuh" + +namespace flashinfer { + +using namespace cute; + +template +struct FP8SparseCollectiveMainloop { + using DTypeQ = typename Ktraits::DTypeQ; + using DTypeKV = typename Ktraits::DTypeKV; + using IdType = typename Ktraits::IdType; + using TileShape_QKD = typename Ktraits::TileShape_QKD; + static constexpr int CTA_Q = get<0>(TileShape_QKD{}); + static constexpr int CTA_KV = get<1>(TileShape_QKD{}); + + static constexpr int NUM_STAGES = Ktraits::NUM_STAGES; + static constexpr int HEAD_DIM = Ktraits::HEAD_DIM; + static constexpr int NUM_MMA_THREADS = Ktraits::NUM_MMA_THREADS; + + using GmemTiledCopyQ = cute::SM90_TMA_LOAD; + static constexpr auto AlignmentKV = 128 / cutlass::sizeof_bits::value; + using AlignmentTypeKV = cute::uint_byte_t(sizeof(DTypeKV)) * AlignmentKV>; + + // Use ZFILL for out-of-bound V loading (avoid nan) + using GmemCopyAtomKV = cute::Copy_Atom, DTypeKV>; + using GmemTiledCopyKV = + decltype(cutlass::gemm::collective::detail::make_simt_gmem_tiled_copy< + GmemCopyAtomKV, Ktraits::NUM_PRODUCER_THREADS, AlignmentKV, + cutlass::detail::TagToStrideB_t, + decltype(cute::get<1>(TileShape_QKD{})), decltype(cute::get<2>(TileShape_QKD{}))>()); + + using SmemLayoutQ = typename Ktraits::SmemLayoutQ; + using SmemLayoutK = typename Ktraits::SmemLayoutK; + using SmemLayoutV = typename Ktraits::SmemLayoutV; + using SmemLayoutVt = typename Ktraits::SmemLayoutVt; + + using ShapeT = cute::Shape; + using StrideT = cute::Shape; // (N, D, H) + using LayoutT = cute::Layout; + + using ShapeLseT = cute::Shape; + using StrideLseT = cute::Shape<_1, int64_t>; + using LayoutLseT = cute::Layout; + + using TMA_Q = decltype(make_tma_copy( + GmemTiledCopyQ{}, + make_tensor(make_gmem_ptr(static_cast(nullptr)), + repeat_like(StrideT{}, int32_t(0)), StrideT{}), + SmemLayoutQ{}, select<0, 2>(TileShape_QKD{}), _1{})); // no mcast for Q + + // for sparse loading, we use cp.async + static constexpr bool USE_TMA_LOAD_KV = false; + using MainloopPipeline = typename Ktraits::MainloopPipeline; + using PipelineParams = typename MainloopPipeline::Params; + using PipelineState = typename MainloopPipeline::PipelineState; + using MainloopPipelineVt = typename Ktraits::MainloopPipelineNoTMA; + using PipelineParamsVt = typename MainloopPipelineVt::Params; + + static constexpr uint32_t TmaTransactionBytesQ = + static_cast(size(SmemLayoutQ{}) * cutlass::sizeof_bits_v / 8); + + static constexpr bool UseSchedulerBarrier = + cutlass::sizeof_bits_v == 8 ? HEAD_DIM >= 128 : HEAD_DIM <= 128; + using WarpScheduler = WarpScheduler; + + // Host side kernel arguments + struct Arguments { + DTypeQ const* Q_ptr; + LayoutT layout_Q; + DTypeKV const* K_ptr; + LayoutT layout_K; + DTypeKV const* V_ptr; + LayoutT layout_V; + IdType const* kv_indices; + int window_left; + AdditionalParams additional_params; + }; + + // Device side kernel params + struct Params { + LayoutT layout_Q; + LayoutT layout_K; + LayoutT layout_V; + TMA_Q tma_load_Q; + DTypeKV* K_ptr; + DTypeKV* V_ptr; + IdType* kv_indices; + int window_left; + AdditionalParams additional_params; + using DTypeKV = typename Ktraits::DTypeKV; + }; + + static Params to_underlying_arguments(Arguments const& args) { + Tensor mQ = make_tensor(make_gmem_ptr(args.Q_ptr), args.layout_Q); + TMA_Q tma_load_Q = + make_tma_copy(GmemTiledCopyQ{}, mQ, SmemLayoutQ{}, select<0, 2>(TileShape_QKD{}), _1{}); + return {args.layout_Q, + args.layout_K, + args.layout_V, + tma_load_Q, + const_cast(args.K_ptr), + const_cast(args.V_ptr), + const_cast(args.kv_indices), + args.window_left, + args.additional_params}; + } + + CUTLASS_DEVICE + static void prefetch_tma_descriptors(Params const& mainloop_params) { + cute::prefetch_tma_descriptor(mainloop_params.tma_load_Q.get_tma_descriptor()); + } + + CUTLASS_DEVICE + int get_num_kv_tiles(Params const& mainloop_params, int q_tile_idx, const int qo_len, + const int kv_len) { + static constexpr int CTA_Q = get<0>(TileShape_QKD{}); + static constexpr int CTA_KV = get<1>(TileShape_QKD{}); + int num_kv_tiles = cute::ceil_div(kv_len, CTA_KV); + if constexpr (CAUSAL) { + num_kv_tiles = std::min(num_kv_tiles, + cute::ceil_div((q_tile_idx + 1) * CTA_Q + kv_len - qo_len, CTA_KV)); + } + + return num_kv_tiles; + } + + template + CUTLASS_DEVICE void load(Params const& mainloop_params, MainloopPipeline pipeline_k, + MainloopPipeline pipeline_v, MainloopPipelineVt pipeline_vt, + PipelineState& smem_pipe_write, PipelineState& smem_pipe_read, + SharedStorage& shared_storage, Scheduler& scheduler, + typename Scheduler::Params const& scheduler_params, + typename Scheduler::WorkTileInfo& work_tile_info, + BlockCoord const& block_coord, int work_idx) { + int thread_idx = threadIdx.x; + int warp_idx_in_warpgroup = __shfl_sync(0xffffffff, (thread_idx / 32) % 4, 0); + bool issue_tma_thread = (warp_idx_in_warpgroup == 0) && (elect_one_sync() == 1); + + Tensor sQ = make_tensor(make_smem_ptr(shared_storage.smem_q.data()), SmemLayoutQ{}); + Tensor sK = make_tensor(make_smem_ptr(shared_storage.smem_k.data()), SmemLayoutK{}); + Tensor sV = make_tensor(make_smem_ptr(shared_storage.smem_v.data()), SmemLayoutV{}); + + Tensor mQ = mainloop_params.tma_load_Q.get_tma_tensor(mainloop_params.layout_Q.shape()); + + // *** Prepare In-kernel V Transpose *** + using SmemLayoutVTransposeSrc = typename Ktraits::SmemLayoutVTransposeSrc; + using SmemLayoutVtTransposeTgt = typename Ktraits::SmemLayoutVtTransposeTgt; + + Tensor sV_src = as_position_independent_swizzle_tensor( + make_tensor(make_smem_ptr(shared_storage.smem_v.data()), SmemLayoutVTransposeSrc{})); + Tensor sVt_tgt = as_position_independent_swizzle_tensor( + make_tensor(make_smem_ptr(shared_storage.smem_vt.data()), SmemLayoutVtTransposeTgt{})); + auto v_tranposer = SmemTransposeFP8_64x64(); + /* ----- V Transpose ---- */ + + auto [q_tile_idx, qo_head_idx, kv_head_idx, qo_indptr, kv_indptr, qo_len, kv_len, batch_idx] = + block_coord; + + // Prepare the TMA loads + Tensor gQ = get_local_tile_tensor(mQ, select<0, 2>(TileShape_QKD{}), qo_head_idx, qo_indptr, + qo_len)(_, _, q_tile_idx); // (Q, D) + + Tensor sQ_x = make_tensor(sQ.data(), make_layout(sQ.layout(), Layout<_1>{})); + Tensor gQ_x = make_tensor(gQ.data(), make_layout(gQ.layout(), Layout<_1>{})); + auto [tQgQ, tQsQ] = + tma_partition(mainloop_params.tma_load_Q, _0{}, Layout<_1>{}, group_modes<0, 2>(sQ_x), + group_modes<0, 2>(gQ_x)); // (TMA), (TMA) + + int num_kv_tiles = get_num_kv_tiles(mainloop_params, q_tile_idx, qo_len, kv_len); + int kv_tile_idx = num_kv_tiles - 1; + int swa_begin_kv_tile_idx = 0; + if constexpr (LEFT_SLIDING_WINDOW) { + swa_begin_kv_tile_idx = get_swa_begin_kv_tile_idx(mainloop_params.window_left, + q_tile_idx, qo_len, kv_len); + } + + constexpr int HEAD_DIM = get<2>(TileShape_QKD{}); + constexpr int CTA_KV = get<1>(TileShape_QKD{}); + auto indexed_gather = BlockSparseIndexedGather(mainloop_params.kv_indices + kv_indptr); + + Tensor mK = make_block_sparse_tensor( // (kv_len, D) + make_gmem_ptr(mainloop_params.K_ptr + kv_head_idx * stride<2>(mainloop_params.layout_K)), + make_shape(kv_len, HEAD_DIM), stride<0>(mainloop_params.layout_K), indexed_gather); + Tensor mV = make_block_sparse_tensor( // (kv_len, D) + make_gmem_ptr(mainloop_params.V_ptr + kv_head_idx * stride<2>(mainloop_params.layout_V)), + make_shape(kv_len, HEAD_DIM), stride<0>(mainloop_params.layout_V), indexed_gather); + + Tensor gK = local_tile(mK, select<1, 2>(TileShape_QKD{}), make_coord(_, _0{})); // (KV, D, kv) + Tensor gV = local_tile(mV, select<1, 2>(TileShape_QKD{}), make_coord(_, _0{})); // (KV, D, kv) + Tensor cKV = cute::make_identity_tensor(gK.shape()); + + GmemTiledCopyKV gmem_tiled_copy_kv; + auto gmem_thr_copy_kv = gmem_tiled_copy_kv.get_slice(thread_idx); + + Tensor tKgK = gmem_thr_copy_kv.partition_S(gK); // (CPY, CPY_KV, CPY_D, kv) + Tensor tKsK = gmem_thr_copy_kv.partition_D(sK); // (CPY, CPY_KV, CPY_D, PIPE) + Tensor tVgV = gmem_thr_copy_kv.partition_S(gV); // (CPY, CPY_KV, CPY_D, kv) + Tensor tVsV = gmem_thr_copy_kv.partition_D(sV); // (CPY, CPY_KV, CPY_D, PIPE) + Tensor tKVcKV = gmem_thr_copy_kv.partition_D(cKV); // (CPY, CPY_KV, CPY_D) + Tensor tKVcKVGroup = flatten_1(tKVcKV); // (CPY, (CPY_KV, CPY_D)) + + int valid_last_kv_tile_size = std::min(kv_len - kv_tile_idx * CTA_KV, CTA_KV); + auto predicate_fn = [&](auto coords) { + auto s_coords = tKVcKVGroup(_0{}, coords); + return elem_less(get<0>(s_coords), valid_last_kv_tile_size); + }; + + // load last k-tile + // all threads are issuing as TMA is disabled + { + pipeline_k.producer_acquire(smem_pipe_write); + Tensor tKgKiGroup = flatten_1(tKgK(_, _, _, kv_tile_idx)); // (CPY, (CPY_KV, CPY_D)) + Tensor tKsKiGroup = + flatten_1(tKsK(_, _, _, smem_pipe_write.index())); // (CPY, (CPY_KV, CPY_D)) + copy_if(gmem_tiled_copy_kv, predicate_fn, tKgKiGroup, tKsKiGroup); + pipeline_k.producer_commit(smem_pipe_write, cutlass::arch::cpasync_barrier_arrive); + } + + // Wait for the MMA warpgroups to say that smem_q is ready + cutlass::arch::NamedBarrier::sync(NUM_MMA_THREADS + Ktraits::NUM_PRODUCER_THREADS, + static_cast(NamedBarriers::kQueryEmpty)); + // load Q tile + if (issue_tma_thread) { + shared_storage.barrier_Q.arrive_and_expect_tx(TmaTransactionBytesQ); + copy(mainloop_params.tma_load_Q.with( + reinterpret_cast( + shared_storage.barrier_Q), + /*mcast_mask=*/0), + tQgQ, tQsQ); + } + + shared_storage.barrier_O.wait((work_idx + 1) % 2); + + if (kv_tile_idx == swa_begin_kv_tile_idx) { + // first tile is the last tile + pipeline_v.producer_acquire(smem_pipe_write); + Tensor tVgViGroup = flatten_1(tVgV(_, _, _, kv_tile_idx)); // (CPY, (CPY_KV, CPY_D)) + Tensor tVsViGroup = + flatten_1(tVsV(_, _, _, smem_pipe_write.index())); // (CPY, (CPY_KV, CPY_D)) + copy_if(gmem_tiled_copy_kv, predicate_fn, tVgViGroup, tVsViGroup); + pipeline_v.producer_commit(smem_pipe_write, cutlass::arch::cpasync_barrier_arrive); + + // Transpose V + pipeline_v.consumer_wait(smem_pipe_read); + pipeline_vt.producer_acquire(smem_pipe_write); + v_tranposer.do_transpose(sV_src, sVt_tgt, smem_pipe_read.index()); + pipeline_vt.producer_commit(smem_pipe_write); // ping MMA consumer + pipeline_v.consumer_release(smem_pipe_read); // release V loading consumer + ++smem_pipe_read; + ++smem_pipe_write; // update state, as K is loaded 1 step faster + } else { + // load second last k-tile and last v-tile + pipeline_v.producer_acquire(smem_pipe_write); + Tensor tVgViGroup = flatten_1(tVgV(_, _, _, kv_tile_idx)); // (CPY, (CPY_KV, CPY_D)) + Tensor tVsViGroup = + flatten_1(tVsV(_, _, _, smem_pipe_write.index())); // (CPY, (CPY_KV, CPY_D)) + copy_if(gmem_tiled_copy_kv, predicate_fn, tVgViGroup, tVsViGroup); + pipeline_v.producer_commit(smem_pipe_write, cutlass::arch::cpasync_barrier_arrive); + + // Transpose V + pipeline_v.consumer_wait(smem_pipe_read); + pipeline_vt.producer_acquire(smem_pipe_write); + v_tranposer.do_transpose(sV_src, sVt_tgt, smem_pipe_read.index()); + pipeline_vt.producer_commit(smem_pipe_write); // ping MMA consumer + pipeline_v.consumer_release(smem_pipe_read); // release V loading consumer + ++smem_pipe_read; + ++smem_pipe_write; // update state, as K is loaded 1 step faster + + pipeline_k.producer_acquire(smem_pipe_write); + Tensor tKgKi = tKgK(_, _, _, kv_tile_idx - 1); // (CPY, CPY_KV, CPY_D) + Tensor tKsKi = tKsK(_, _, _, smem_pipe_write.index()); // (CPY, CPY_KV, CPY_D) + copy(gmem_tiled_copy_kv, tKgKi, tKsKi); + pipeline_k.producer_commit(smem_pipe_write, cutlass::arch::cpasync_barrier_arrive); + + --kv_tile_idx; + + // load remaining k/v tiles +#pragma unroll 2 + for (; kv_tile_idx > swa_begin_kv_tile_idx; --kv_tile_idx) { + pipeline_v.producer_acquire(smem_pipe_write); + Tensor tVgVi = tVgV(_, _, _, kv_tile_idx); // (CPY, CPY_KV, CPY_D) + Tensor tVsVi = tVsV(_, _, _, smem_pipe_write.index()); // (CPY, CPY_KV, CPY_D) + copy(gmem_tiled_copy_kv, tVgVi, tVsVi); + pipeline_v.producer_commit(smem_pipe_write, cutlass::arch::cpasync_barrier_arrive); + + // Transpose V + pipeline_v.consumer_wait(smem_pipe_read); + pipeline_vt.producer_acquire(smem_pipe_write); + v_tranposer.do_transpose(sV_src, sVt_tgt, smem_pipe_read.index()); + pipeline_vt.producer_commit(smem_pipe_write); // ping MMA consumer + pipeline_v.consumer_release(smem_pipe_read); // release V loading consumer + ++smem_pipe_read; + ++smem_pipe_write; // update state, as K is loaded 1 step faster + + pipeline_k.producer_acquire(smem_pipe_write); + Tensor tKgKi = tKgK(_, _, _, kv_tile_idx - 1); // (CPY, CPY_KV, CPY_D) + Tensor tKsKi = tKsK(_, _, _, smem_pipe_write.index()); // (CPY, CPY_KV, CPY_D) + copy(gmem_tiled_copy_kv, tKgKi, tKsKi); + pipeline_k.producer_commit(smem_pipe_write, cutlass::arch::cpasync_barrier_arrive); + } + scheduler.prefetch_next_work(scheduler_params, work_tile_info); + + // load first v tile + { + pipeline_v.producer_acquire(smem_pipe_write); + Tensor tVgVi = tVgV(_, _, _, 0); // (CPY, (CPY_KV, CPY_D)) + Tensor tVsVi = tVsV(_, _, _, smem_pipe_write.index()); // (CPY, (CPY_KV, CPY_D)) + copy(gmem_tiled_copy_kv, tVgVi, tVsVi); + pipeline_v.producer_commit(smem_pipe_write, cutlass::arch::cpasync_barrier_arrive); + + // Transpose V + pipeline_v.consumer_wait(smem_pipe_read); + pipeline_vt.producer_acquire(smem_pipe_write); + v_tranposer.do_transpose(sV_src, sVt_tgt, smem_pipe_read.index()); + pipeline_vt.producer_commit(smem_pipe_write); // ping MMA consumer + pipeline_v.consumer_release(smem_pipe_read); // release V loading consumer + ++smem_pipe_read; + ++smem_pipe_write; // update state, as K is loaded 1 step faster + } + } + + scheduler.broadcast_next_work(work_tile_info); + } + + CUTLASS_DEVICE void load_tail(MainloopPipeline pipeline_k, MainloopPipeline pipeline_v, + PipelineState& smem_pipe_write) { + pipeline_k.producer_tail(smem_pipe_write); + pipeline_v.producer_tail(smem_pipe_write); + } +}; + +} // namespace flashinfer + +#endif // FLASHINFER_ATTENTION_HOPPER_FP8_SPARSE_MAINLOOP_CUH_ diff --git a/include/flashinfer/attention/hopper/quantization/prefill_sm90.cuh b/include/flashinfer/attention/hopper/quantization/prefill_sm90.cuh new file mode 100644 index 0000000000000000000000000000000000000000..24e416b61bf9b95795144f19362238b42ca8ae31 --- /dev/null +++ b/include/flashinfer/attention/hopper/quantization/prefill_sm90.cuh @@ -0,0 +1,477 @@ +/* + * Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri + * Dao. Licensed under the BSD 3-Clause. + * + * Modified by the FlashInfer team. + */ +#ifndef FLASHINFER_ATTENTION_HOPPER_FP8_PREFILL_SM90_CUH_ +#define FLASHINFER_ATTENTION_HOPPER_FP8_PREFILL_SM90_CUH_ + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +#include "../../../cutlass_utils.cuh" +#include "../../../exception.h" +#include "../../mask.cuh" +#include "../sparse_mainloop.cuh" +#include "../tile_scheduler.cuh" +#include "../utils.cuh" +#include "epilogue.cuh" +#include "kernel_traits.cuh" +#include "mainloop_load.cuh" +#include "mainloop_mma.cuh" +#include "mainloop_sparse_load.cuh" + +namespace flashinfer { + +using namespace cute; + +template +__global__ void __launch_bounds__(Ktraits::NUM_WARPS* cutlass::NumThreadsPerWarp, 1) + FP8PrefillWithKVCacheKernel(CUTE_GRID_CONSTANT + typename CollectiveMainloop::Params const mainloop_params, + CUTE_GRID_CONSTANT + typename CollectiveEpilogue::Params const epilogue_params, + CUTE_GRID_CONSTANT + typename TileScheduler::Params const scheduler_params) { + using DTypeQ = typename Ktraits::DTypeQ; + using DTypeKV = typename Ktraits::DTypeKV; + using DTypeO = typename Ktraits::DTypeO; + using DTypeQKAccum = typename Ktraits::DTypeQKAccum; + using TileShape_QKD = typename Ktraits::TileShape_QKD; + using AttentionVariant = typename Ktraits::AttentionVariant; + + static constexpr int NUM_MMA_THREADS = Ktraits::NUM_MMA_THREADS; + // We always assign one WG as producer + // For FP8 kernel, all 4 warps collectively process ldmatrix with ldmatrix + static constexpr int NUM_COPY_THREADS = Ktraits::NUM_PRODUCER_THREADS; + static constexpr int CTA_Q = Ktraits::CTA_Q; + static constexpr int CTA_KV = Ktraits::CTA_KV; + + static constexpr bool use_tma_load_kv = CollectiveMainloop::USE_TMA_LOAD_KV; + // Pipeline for loading K/V + using MainloopPipeline = typename CollectiveMainloop::MainloopPipeline; + using PipelineParams = typename MainloopPipeline::Params; + using PipelineState = typename MainloopPipeline::PipelineState; + + // Pipeline for transposing V + using MainloopPipelineVt = typename CollectiveMainloop::MainloopPipelineVt; + using PipelineParamsVt = typename MainloopPipelineVt::Params; + + extern __shared__ char shared_memory[]; + auto& shared_storage = *reinterpret_cast(shared_memory); + + int const lane_predicate = cute::elect_one_sync(); + int const warp_idx = cutlass::canonical_warp_idx_sync(); + + // Issue Tma Descriptor Prefetch from a single thread + if (warp_idx == 0 && lane_predicate) { + CollectiveMainloop::prefetch_tma_descriptors(mainloop_params); + CollectiveEpilogue::prefetch_tma_descriptors(epilogue_params); + } + + // Obtain warp index + int const warp_group_thread_idx = threadIdx.x % cutlass::NumThreadsPerWarpGroup; + + PipelineParams pipeline_params; + int warp_group_idx = cutlass::canonical_warp_group_idx(); + pipeline_params.role = warp_group_idx == 0 ? MainloopPipeline::ThreadCategory::Producer + : MainloopPipeline::ThreadCategory::Consumer; + if constexpr (use_tma_load_kv) { + pipeline_params.transaction_bytes = CollectiveMainloop::TmaTransactionBytesK; + pipeline_params.is_leader = warp_group_thread_idx == 0; + pipeline_params.num_consumers = NUM_MMA_THREADS; + } else { + pipeline_params.producer_arv_count = NUM_COPY_THREADS; + pipeline_params.consumer_arv_count = NUM_MMA_THREADS; + } + + if (warp_idx == 0 && lane_predicate) { + shared_storage.barrier_Q.init(/*num_threads=*/1); + shared_storage.barrier_O.init(/*num_threads=*/1); + } + // We're counting on pipeline_k to call cutlass::arch::fence_barrier_init(); + MainloopPipeline pipeline_k = [&] { + if constexpr (use_tma_load_kv) { + return MainloopPipeline(shared_storage.pipeline_k, pipeline_params, + /*cluster_shape=*/Shape<_1, _1, _1>{}); + } else { + return MainloopPipeline(shared_storage.pipeline_k, pipeline_params); + } + }(); + + MainloopPipeline pipeline_v = [&] { + // specialized for shared memory of V transpose + pipeline_params.role = MainloopPipeline::ThreadCategory::ProducerConsumer; + if constexpr (use_tma_load_kv) { + pipeline_params.num_consumers = NUM_COPY_THREADS; + return MainloopPipeline(shared_storage.pipeline_v, pipeline_params, + /*cluster_shape=*/Shape<_1, _1, _1>{}); + } else { + pipeline_params.consumer_arv_count = NUM_COPY_THREADS; + return MainloopPipeline(shared_storage.pipeline_v, pipeline_params); + } + }(); + + // Init pipeline_vt for transpose and consumed by mma + PipelineParamsVt pipeline_params_vt; + pipeline_params_vt.producer_arv_count = NUM_COPY_THREADS; + pipeline_params_vt.consumer_arv_count = NUM_MMA_THREADS; + MainloopPipelineVt pipeline_vt(shared_storage.pipeline_vt, pipeline_params_vt); + + CollectiveMainloop collective_mainloop; + CollectiveEpilogue collective_epilogue; + + // We need this to guarantee that the Pipeline init is visible to all producers and consumer + // blocks in the Cluster + __syncthreads(); + + if (warp_group_idx == 0) { // Producer + if constexpr (use_tma_load_kv) { + cutlass::arch::warpgroup_reg_dealloc(); + } else { + cutlass::arch::warpgroup_reg_dealloc<72>(); + } + + // Here no condition as the entire warp group is used as producer + PipelineState smem_pipe_write = cutlass::make_producer_start_state(); + PipelineState smem_pipe_read; + + int work_idx = 0; + + TileScheduler scheduler; + for (auto work_tile_info = scheduler.get_initial_work(scheduler_params); + work_tile_info.is_valid(scheduler_params); + work_tile_info = scheduler.template get_next_work(scheduler_params, + work_tile_info)) { + auto block_coord = work_tile_info.get_block_coord(scheduler_params); + auto [q_tile_idx, qo_head_idx, kv_head_idx, qo_indptr, kv_indptr, qo_len, kv_len, batch_idx] = + block_coord; + + if (q_tile_idx * CTA_Q >= qo_len) { + continue; + } + int num_kv_tiles = + collective_mainloop.get_num_kv_tiles(mainloop_params, q_tile_idx, qo_len, kv_len); + if (num_kv_tiles <= 0) { + scheduler.prefetch_next_work(scheduler_params, work_tile_info); + scheduler.broadcast_next_work(work_tile_info); + continue; + } + collective_mainloop.load( + mainloop_params, pipeline_k, pipeline_v, pipeline_vt, smem_pipe_write, smem_pipe_read, + shared_storage, scheduler, scheduler_params, work_tile_info, block_coord, work_idx); + ++work_idx; + } + collective_mainloop.load_tail(pipeline_k, pipeline_v, smem_pipe_write); + + } else { // Consumer + if constexpr (use_tma_load_kv) { + cutlass::arch::warpgroup_reg_alloc(); + } else { + cutlass::arch::warpgroup_reg_alloc(); + } + + TileScheduler scheduler; + // Initialize matmul objects. + typename Ktraits::TiledMmaPV tiled_mma_pv; + + PipelineState smem_pipe_read_k, smem_pipe_read_v; + // We don't need separate variables smem_pipe_release_k and smem_pipe_release_v + // (like in Cutlass's gemm) because the read and release pipeline states are always the same. + + CollectiveMainloop::WarpScheduler::mma_init(); + scheduler.init_consumer(); + + int work_idx = 0; + CUTLASS_PRAGMA_NO_UNROLL + for (auto work_tile_info = scheduler.get_initial_work(scheduler_params); + work_tile_info.is_valid(scheduler_params); + work_tile_info = scheduler.template get_next_work(scheduler_params, + work_tile_info)) { + // Attention output (GEMM-II) accumulator. + Tensor tOrO = partition_fragment_C(tiled_mma_pv, select<0, 2>(TileShape_QKD{})); + clear(tOrO); + + auto block_coord = work_tile_info.get_block_coord(scheduler_params); + auto [q_tile_idx, qo_head_idx, kv_head_idx, qo_indptr, kv_indptr, qo_len, kv_len, batch_idx] = + block_coord; + + AttentionVariant variant(mainloop_params, block_coord); + auto attention_updater = + variant.template GetAttentionUpdater<2 * (2 * CTA_Q / NUM_MMA_THREADS)>(); + + if (q_tile_idx * CTA_Q >= qo_len) { + continue; + } + int num_kv_tiles = + collective_mainloop.get_num_kv_tiles(mainloop_params, q_tile_idx, qo_len, kv_len); + if (num_kv_tiles <= 0) { // We exit early and write 0 to gO and -inf to gLSE. + collective_epilogue.store_zero(epilogue_params, shared_storage, + threadIdx.x - NUM_COPY_THREADS, block_coord); + continue; + } + + int swa_begin_kv_tile_idx = 0; + int swa_end_kv_tile_idx = -1; + if constexpr (LEFT_SLIDING_WINDOW) { + swa_begin_kv_tile_idx = get_swa_begin_kv_tile_idx( + mainloop_params.window_left, q_tile_idx, qo_len, kv_len); + swa_end_kv_tile_idx = get_swa_end_kv_tile_idx(mainloop_params.window_left, + q_tile_idx, qo_len, kv_len); + } + + mma_fp8( + mainloop_params, variant, pipeline_k, pipeline_vt, smem_pipe_read_k, smem_pipe_read_v, + tOrO, attention_updater, num_kv_tiles, swa_begin_kv_tile_idx, swa_end_kv_tile_idx, + threadIdx.x - NUM_COPY_THREADS, work_idx, q_tile_idx, shared_storage, qo_len, kv_len, + qo_head_idx, kv_head_idx, batch_idx); + + collective_epilogue.store(epilogue_params, tOrO, attention_updater.get_lse(), shared_storage, + tiled_mma_pv, threadIdx.x - NUM_COPY_THREADS, block_coord); + + ++work_idx; + } + collective_epilogue.store_tail(); + } +} + +template +cudaError_t SingleFP8PrefillWithKVCacheKernelTraitsDispatched(Params& params, cudaStream_t stream) { + using DTypeQ = typename KernelTraits::DTypeQ; + using DTypeKV = typename KernelTraits::DTypeKV; + using DTypeO = typename KernelTraits::DTypeO; + using TileShape_QKD = typename KernelTraits::TileShape_QKD; + + using CollectiveMainloop = + FP8CollectiveMainloop; + using CollectiveEpilogue = FP8CollectiveEpilogue; + using Scheduler = SingleTileScheduler; + typename CollectiveMainloop::Params mainloop_params = CollectiveMainloop::to_underlying_arguments( + {params.q_ptr, + get_gmem_layout(params.qo_len, params.num_qo_heads, KernelTraits::HEAD_DIM, + params.q_stride_n, + params.q_stride_h), // layout_Q + params.k_ptr, + get_gmem_layout(params.kv_len, params.num_kv_heads, KernelTraits::HEAD_DIM, + params.k_stride_n, + params.k_stride_h), // layout_K + params.v_ptr, + get_gmem_layout(params.kv_len, params.num_kv_heads, KernelTraits::HEAD_DIM, + params.v_stride_n, + params.v_stride_h), // layout_V + params.window_left, params.additional_params}); + typename CollectiveEpilogue::Params epilogue_params = + CollectiveEpilogue::to_underlying_arguments({ + static_cast(params.o_ptr), + get_gmem_layout(params.qo_len, params.num_qo_heads, KernelTraits::HEAD_DIM, + params.o_stride_n, + params.o_stride_h), // layout_O + static_cast(params.lse_ptr), + get_lse_gmem_layout(params.qo_len, params.num_qo_heads), // layout_LSE + }); + + int num_tiles_q = cutlass::ceil_div(params.qo_len, KernelTraits::CTA_Q); + // TODO(Zihao): also support kv-head major + typename Scheduler::Arguments scheduler_args = { + num_tiles_q, params.num_qo_heads, params.qo_len, params.kv_len, + cutlass::FastDivmod(params.num_qo_heads / params.num_kv_heads)}; + typename Scheduler::Params scheduler_params = Scheduler::to_underlying_arguments(scheduler_args); + + auto kernel = + (void*)FP8PrefillWithKVCacheKernel; + int smem_size = sizeof(typename KernelTraits::SharedStorage); + FLASHINFER_CUDA_CALL( + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + + int device; + cudaGetDevice(&device); + int multiprocessor_count; + FLASHINFER_CUDA_CALL( + cudaDeviceGetAttribute(&multiprocessor_count, cudaDevAttrMultiProcessorCount, device)); + dim3 grid_dims = Scheduler::get_grid_dim(scheduler_args, multiprocessor_count); + static constexpr int num_ctas = KernelTraits::NUM_WARPS * 32; + dim3 block_dims(num_ctas); + void* args[] = {&mainloop_params, &epilogue_params, &scheduler_params}; + FLASHINFER_CUDA_CALL(cudaLaunchKernel(kernel, grid_dims, block_dims, args, smem_size, stream)); + + return cudaSuccess; +} + +template +cudaError_t BatchFP8PrefillWithPagedKVCacheKernelTraitsDispatched(Params& params, + cudaStream_t stream) { + using DTypeQ = typename KernelTraits::DTypeQ; + using DTypeKV = typename KernelTraits::DTypeKV; + using DTypeO = typename KernelTraits::DTypeO; + using IdType = typename KernelTraits::IdType; + + using CollectiveMainloop = + FP8SparseCollectiveMainloop; + using CollectiveEpilogue = FP8CollectiveEpilogue; + using Scheduler = + std::conditional_t, + BatchPrefillPersistentTileScheduler>; + + typename CollectiveMainloop::Params mainloop_params = CollectiveMainloop::to_underlying_arguments( + {params.q_ptr, + get_gmem_layout(params.nnz_qo, params.num_qo_heads, KernelTraits::HEAD_DIM, + params.q_stride_n, + params.q_stride_h), // layout_Q + params.k_ptr, + // NOTE(Zihao): nnz was useless here, we can just pass 0 + get_gmem_layout(/*nnz=*/0, params.num_kv_heads, KernelTraits::HEAD_DIM, params.k_stride_n, + params.k_stride_h), // layout_K + params.v_ptr, + get_gmem_layout(/*nnz=*/0, params.num_kv_heads, KernelTraits::HEAD_DIM, params.v_stride_n, + params.v_stride_h), // layout_V + params.kv_indices, params.window_left, params.additional_params}); + typename CollectiveEpilogue::Params epilogue_params = + CollectiveEpilogue::to_underlying_arguments({ + params.o_ptr, + get_gmem_layout(params.nnz_qo, params.num_qo_heads, KernelTraits::HEAD_DIM, + params.o_stride_n, + params.o_stride_h), // layout_O + params.lse_ptr, get_lse_gmem_layout(params.nnz_qo, params.num_qo_heads), // layout_LSE + }); + + typename Scheduler::Arguments scheduler_args = { + params.work_indptr, + params.head_indices, + params.qo_tile_indices, + params.qo_indptr, + params.kv_indptr, + params.qo_lens, + params.kv_lens, + params.batch_indices, + cutlass::FastDivmod(params.num_qo_heads / params.num_kv_heads), + params.num_qo_heads}; + typename Scheduler::Params scheduler_params = Scheduler::to_underlying_arguments(scheduler_args); + + // Get the ptr to kernel function. + auto kernel = + (void*)FP8PrefillWithKVCacheKernel; + int smem_size = sizeof(typename KernelTraits::SharedStorage); + FLASHINFER_CUDA_CALL( + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + + int device; + cudaGetDevice(&device); + int multiprocessor_count; + FLASHINFER_CUDA_CALL( + cudaDeviceGetAttribute(&multiprocessor_count, cudaDevAttrMultiProcessorCount, device)); + dim3 grid_dims = Scheduler::get_grid_dim(scheduler_args, multiprocessor_count); + static constexpr int ctaSize = KernelTraits::NUM_WARPS * 32; + dim3 block_dims(ctaSize); + void* args[] = {&mainloop_params, &epilogue_params, &scheduler_params}; + FLASHINFER_CUDA_CALL(cudaLaunchKernel(kernel, grid_dims, block_dims, args, smem_size, stream)); + + return cudaSuccess; +} + +template +cudaError_t SingleFP8PrefillWithKVCacheDispatched(Params& params, cudaStream_t stream) { + static_assert(cutlass::sizeof_bits_v == 8); + static_assert(cutlass::sizeof_bits_v == 8); + static_assert(HEAD_DIM == 64 || HEAD_DIM == 128 || HEAD_DIM == 256); + if (MASK_MODE == MaskMode::kCustom) { + return cudaErrorNotSupported; // Not supported yet. + } + constexpr bool CAUSAL = MASK_MODE == MaskMode::kCausal; + if constexpr (HEAD_DIM == 64) { + SingleFP8PrefillWithKVCacheKernelTraitsDispatched< + FP8AttentionKernelTraits, + LEFT_SLIDING_WINDOW, CAUSAL>(params, stream); + } else if constexpr (HEAD_DIM == 128) { + SingleFP8PrefillWithKVCacheKernelTraitsDispatched< + FP8AttentionKernelTraits, + LEFT_SLIDING_WINDOW, CAUSAL>(params, stream); + } else { + // HEAD_DIM == 256; + SingleFP8PrefillWithKVCacheKernelTraitsDispatched< + FP8AttentionKernelTraits, + LEFT_SLIDING_WINDOW, CAUSAL>(params, stream); + } + cudaError_t status = cudaGetLastError(); + return status; +} + +template +cudaError_t BatchFP8PrefillWithPagedKVCacheDispatched(Params& params, bool enable_pdl, + cudaStream_t stream) { + static_assert(HEAD_DIM == 64 || HEAD_DIM == 128 || HEAD_DIM == 256); + if (MASK_MODE == MaskMode::kCustom) { + return cudaErrorNotSupported; // Not supported yet. + } + constexpr bool CAUSAL = MASK_MODE == MaskMode::kCausal; + if constexpr (HEAD_DIM == 64) { + // NOTE(Zihao): CTA_KV not tuned for HEAD_DIM == 64, need to optimize later + BatchFP8PrefillWithPagedKVCacheKernelTraitsDispatched< + FP8AttentionKernelTraits, + LEFT_SLIDING_WINDOW, CAUSAL, SAME_SCHEDULE_FOR_ALL_HEADS>(params, stream); + } else if constexpr (HEAD_DIM == 128) { + BatchFP8PrefillWithPagedKVCacheKernelTraitsDispatched< + FP8AttentionKernelTraits, + LEFT_SLIDING_WINDOW, CAUSAL, SAME_SCHEDULE_FOR_ALL_HEADS>(params, stream); + } else { + // HEAD_DIM == 256; + // NOTE(Zihao): CTA_KV not tuned for HEAD_DIM == 256, need to optimize later + BatchFP8PrefillWithPagedKVCacheKernelTraitsDispatched< + FP8AttentionKernelTraits, + LEFT_SLIDING_WINDOW, CAUSAL, SAME_SCHEDULE_FOR_ALL_HEADS>(params, stream); + } + cudaError_t status = cudaGetLastError(); + return status; +}; + +} // namespace flashinfer + +#endif // FLASHINFER_ATTENTION_HOPPER_FP8_PREFILL_SM90_CUH_ diff --git a/include/flashinfer/attention/hopper/sparse_mainloop.cuh b/include/flashinfer/attention/hopper/sparse_mainloop.cuh new file mode 100644 index 0000000000000000000000000000000000000000..cc713cdfb2bf4b1f1583316a1ea1d663183c5e2a --- /dev/null +++ b/include/flashinfer/attention/hopper/sparse_mainloop.cuh @@ -0,0 +1,362 @@ +/* + * Copyright (c) 2024 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef FLASHINFER_ATTENTION_HOPPER_SPARSE_MAINLOOP_CUH_ +#define FLASHINFER_ATTENTION_HOPPER_SPARSE_MAINLOOP_CUH_ + +#include +#include +#include +#include + +#include "../../math.cuh" +#include "block_sparse_gather.cuh" +#include "cute/tensor.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/pipeline/pipeline.hpp" +#include "named_barrier.cuh" +#include "utils.cuh" + +namespace flashinfer { + +using namespace cute; + +template +struct SparseCollectiveMainloop { + using DTypeQ = typename Ktraits::DTypeQ; + using DTypeKV = typename Ktraits::DTypeKV; + using IdType = typename Ktraits::IdType; + using TileShape_QKD = typename Ktraits::TileShape_QKD; + using TileShape_PDV = typename Ktraits::TileShape_PDV; + static constexpr int CTA_Q = get<0>(TileShape_QKD{}); + static constexpr int CTA_KV = get<1>(TileShape_QKD{}); + + static constexpr int NUM_STAGES = Ktraits::NUM_STAGES; + static constexpr int HEAD_DIM_QK = Ktraits::HEAD_DIM_QK; + static constexpr int HEAD_DIM_VO = Ktraits::HEAD_DIM_VO; + static_assert(HEAD_DIM_QK == HEAD_DIM_VO); + static constexpr int NUM_COPY_THREADS = cutlass::NumThreadsPerWarpGroup; + + using GmemTiledCopyQ = cute::SM90_TMA_LOAD; + static constexpr auto AlignmentKV = 128 / cutlass::sizeof_bits::value; + using AlignmentTypeKV = cute::uint_byte_t(sizeof(DTypeKV)) * AlignmentKV>; + // NOTE(Zihao): use SM80_CP_ASYNC for sparse loading of KV-cache + using GmemCopyAtomKV = cute::Copy_Atom, DTypeKV>; + using GmemTiledCopyK = + decltype(cutlass::gemm::collective::detail::make_simt_gmem_tiled_copy< + GmemCopyAtomKV, NUM_COPY_THREADS, AlignmentKV, + cutlass::detail::TagToStrideB_t, + decltype(cute::get<1>(TileShape_QKD{})), decltype(cute::get<2>(TileShape_QKD{}))>()); + using GmemTiledCopyV = + decltype(cutlass::gemm::collective::detail::make_simt_gmem_tiled_copy< + GmemCopyAtomKV, NUM_COPY_THREADS, AlignmentKV, + cutlass::detail::TagToStrideB_t, + decltype(cute::get<2>(TileShape_PDV{})), decltype(cute::get<1>(TileShape_PDV{}))>()); + + using SmemLayoutQ = typename Ktraits::SmemLayoutQ; + using SmemLayoutK = typename Ktraits::SmemLayoutK; + using SmemLayoutV = typename Ktraits::SmemLayoutV; + using SmemLayoutVt = typename Ktraits::SmemLayoutVt; + + using ShapeT = cute::Shape; + using StrideT = cute::Shape; // (N, D, H) + using LayoutT = cute::Layout; + + using ShapeLseT = cute::Shape; + using StrideLseT = cute::Shape<_1, int64_t>; + using LayoutLseT = cute::Layout; + + using TMA_Q = decltype(make_tma_copy( + GmemTiledCopyQ{}, + make_tensor(make_gmem_ptr(static_cast(nullptr)), + repeat_like(StrideT{}, int32_t(0)), StrideT{}), + SmemLayoutQ{}, select<0, 2>(TileShape_QKD{}), _1{})); // no mcast for Q + + static constexpr bool USE_TMA_LOAD_KV = false; + static constexpr int NUM_MMA_THREADS = size(typename Ktraits::TiledMmaQK{}); + using MainloopPipeline = typename Ktraits::MainloopPipeline; + using PipelineParams = typename MainloopPipeline::Params; + using PipelineState = typename MainloopPipeline::PipelineState; + + static constexpr uint32_t TmaTransactionBytesQ = + static_cast(size(SmemLayoutQ{}) * cutlass::sizeof_bits_v / 8); + + static constexpr bool UseSchedulerBarrier = + cutlass::sizeof_bits_v == 8 ? HEAD_DIM_VO >= 128 : HEAD_DIM_VO <= 128; + using WarpScheduler = WarpScheduler; + + // Host side kernel arguments + struct Arguments { + DTypeQ const* Q_ptr; + LayoutT layout_Q; + DTypeKV const* K_ptr; + LayoutT layout_K; + DTypeKV const* V_ptr; + LayoutT layout_V; + IdType const* kv_indices; + int window_left; + AdditionalParams additional_params; + }; + + // Device side kernel params + struct Params { + LayoutT layout_Q; + LayoutT layout_K; + LayoutT layout_V; + TMA_Q tma_load_Q; + DTypeKV* K_ptr; + DTypeKV* V_ptr; + IdType* kv_indices; + int window_left; + AdditionalParams additional_params; + }; + + static Params to_underlying_arguments(Arguments const& args) { + Tensor mQ = make_tensor(make_gmem_ptr(args.Q_ptr), args.layout_Q); + TMA_Q tma_load_Q = + make_tma_copy(GmemTiledCopyQ{}, mQ, SmemLayoutQ{}, select<0, 2>(TileShape_QKD{}), _1{}); + return {args.layout_Q, + args.layout_K, + args.layout_V, + tma_load_Q, + const_cast(args.K_ptr), + const_cast(args.V_ptr), + const_cast(args.kv_indices), + args.window_left, + args.additional_params}; + } + + CUTLASS_DEVICE + static void prefetch_tma_descriptors(Params const& mainloop_params) { + cute::prefetch_tma_descriptor(mainloop_params.tma_load_Q.get_tma_descriptor()); + } + + CUTLASS_DEVICE + int get_num_kv_tiles(Params const& mainloop_params, int q_tile_idx, const int qo_len, + const int kv_len) { + static constexpr int CTA_Q = get<0>(TileShape_QKD{}); + static constexpr int CTA_KV = get<1>(TileShape_QKD{}); + int num_kv_tiles = cute::ceil_div(kv_len, CTA_KV); + if constexpr (CAUSAL) { + num_kv_tiles = std::min(num_kv_tiles, + cute::ceil_div((q_tile_idx + 1) * CTA_Q + kv_len - qo_len, CTA_KV)); + } + if constexpr (MULTIITEMSCORING) { + num_kv_tiles = std::min(num_kv_tiles, + cute::ceil_div((q_tile_idx + 1) * CTA_Q + kv_len - qo_len, CTA_KV)); + } + + return num_kv_tiles; + } + + template + CUTLASS_DEVICE void load(Params const& mainloop_params, MainloopPipeline pipeline_k, + MainloopPipeline pipeline_v, PipelineState& smem_pipe_write_k, + PipelineState& smem_pipe_write_v, SharedStorage& shared_storage, + Scheduler& scheduler, typename Scheduler::Params const& scheduler_params, + typename Scheduler::WorkTileInfo& work_tile_info, + BlockCoord const& block_coord, int work_idx, + const int num_kv_tiles_outside_items_window = 0, + const int num_kv_tiles_prefix = 0) { + int thread_idx = threadIdx.x; + int warp_idx_in_warpgroup = __shfl_sync(0xffffffff, (thread_idx / 32) % 4, 0); + Tensor sQ = make_tensor(make_smem_ptr(shared_storage.smem_q.data()), SmemLayoutQ{}); + Tensor sK = make_tensor(make_smem_ptr(shared_storage.smem_k.data()), SmemLayoutK{}); + Tensor sV = make_tensor(make_smem_ptr(shared_storage.smem_v.data()), SmemLayoutV{}); + + Tensor mQ = mainloop_params.tma_load_Q.get_tma_tensor(mainloop_params.layout_Q.shape()); + + auto [q_tile_idx, qo_head_idx, kv_head_idx, qo_indptr, kv_indptr, qo_len, kv_len, batch_idx] = + block_coord; + + // Prepare the TMA loads + Tensor gQ = get_local_tile_tensor(mQ, select<0, 2>(TileShape_QKD{}), qo_head_idx, qo_indptr, + qo_len)(_, _, q_tile_idx); // (Q, D) + + Tensor sQ_x = make_tensor(sQ.data(), make_layout(sQ.layout(), Layout<_1>{})); + Tensor gQ_x = make_tensor(gQ.data(), make_layout(gQ.layout(), Layout<_1>{})); + auto [tQgQ, tQsQ] = + tma_partition(mainloop_params.tma_load_Q, _0{}, Layout<_1>{}, group_modes<0, 2>(sQ_x), + group_modes<0, 2>(gQ_x)); // (TMA), (TMA) + + int num_kv_tiles = get_num_kv_tiles(mainloop_params, q_tile_idx, qo_len, kv_len); + int kv_tile_idx = num_kv_tiles - 1; + int swa_begin_kv_tile_idx = 0; + if constexpr (LEFT_SLIDING_WINDOW) { + swa_begin_kv_tile_idx = get_swa_begin_kv_tile_idx(mainloop_params.window_left, + q_tile_idx, qo_len, kv_len); + } + + constexpr int HEAD_DIM_QK = get<2>(TileShape_QKD{}); + constexpr int HEAD_DIM_VO = get<1>(TileShape_PDV{}); + constexpr int CTA_KV = get<1>(TileShape_QKD{}); + auto indexed_gather = BlockSparseIndexedGather(mainloop_params.kv_indices + kv_indptr); + + Tensor mK = make_block_sparse_tensor( // (kv_len, D_K) + make_gmem_ptr(mainloop_params.K_ptr + kv_head_idx * stride<2>(mainloop_params.layout_K)), + make_shape(kv_len, HEAD_DIM_QK), stride<0>(mainloop_params.layout_K), indexed_gather); + Tensor mV = make_block_sparse_tensor( // (kv_len, D_V) + make_gmem_ptr(mainloop_params.V_ptr + kv_head_idx * stride<2>(mainloop_params.layout_V)), + make_shape(kv_len, HEAD_DIM_VO), stride<0>(mainloop_params.layout_V), indexed_gather); + + Tensor gK = + local_tile(mK, select<1, 2>(TileShape_QKD{}), make_coord(_, _0{})); // (KV, D_K, kv) + Tensor gV = + local_tile(mV, select<2, 1>(TileShape_PDV{}), make_coord(_, _0{})); // (KV, D_V, kv) + Tensor cK = cute::make_identity_tensor(gK.shape()); + Tensor cV = cute::make_identity_tensor(gV.shape()); + + GmemTiledCopyK gmem_tiled_copy_k; + GmemTiledCopyV gmem_tiled_copy_v; + auto gmem_thr_copy_k = gmem_tiled_copy_k.get_slice(thread_idx); + auto gmem_thr_copy_v = gmem_tiled_copy_v.get_slice(thread_idx); + + Tensor tKgK = gmem_thr_copy_k.partition_S(gK); // (CPY, CPY_KV, CPY_D, kv) + Tensor tKsK = gmem_thr_copy_k.partition_D(sK); // (CPY, CPY_KV, CPY_D, PIPE) + Tensor tVgV = gmem_thr_copy_v.partition_S(gV); // (CPY, CPY_KV, CPY_D, kv) + Tensor tVsV = gmem_thr_copy_v.partition_D(sV); // (CPY, CPY_KV, CPY_D, PIPE) + Tensor tKcK = gmem_thr_copy_k.partition_D(cK); // (CPY, CPY_KV, CPY_D) + Tensor tKcKGroup = flatten_1(tKcK); // (CPY, (CPY_KV, CPY_D)) + Tensor tVcV = gmem_thr_copy_v.partition_D(cV); // (CPY, CPY_KV, CPY_D) + Tensor tVcVGroup = flatten_1(tVcV); // (CPY, (CPY_KV, CPY_D)) + + int valid_last_kv_tile_size = std::min(kv_len - kv_tile_idx * CTA_KV, CTA_KV); + auto k_predicate_fn = [&](auto coords) { + auto s_coords = tKcKGroup(_0{}, coords); + return elem_less(get<0>(s_coords), valid_last_kv_tile_size); + }; + auto v_predicate_fn = [&](auto coords) { + auto s_coords = tVcVGroup(_0{}, coords); + return elem_less(get<0>(s_coords), valid_last_kv_tile_size); + }; + auto kv_tile_idx_decrement = [&](int kv_tile_idx) { + int result = kv_tile_idx - 1; + if constexpr (MULTIITEMSCORING) { + if ((kv_tile_idx == num_kv_tiles_outside_items_window) & + (kv_tile_idx >= num_kv_tiles_prefix)) { + result = num_kv_tiles_prefix - 1; + } + } + return result; + }; + + // load last k-tile + { + pipeline_k.producer_acquire(smem_pipe_write_k); + Tensor tKgKiGroup = flatten_1(tKgK(_, _, _, kv_tile_idx)); // (CPY, (CPY_KV, CPY_D)) + Tensor tKsKiGroup = + flatten_1(tKsK(_, _, _, smem_pipe_write_k.index())); // (CPY, (CPY_KV, CPY_D)) + copy_if(gmem_tiled_copy_k, k_predicate_fn, tKgKiGroup, tKsKiGroup); + + pipeline_k.producer_commit(smem_pipe_write_k, cutlass::arch::cpasync_barrier_arrive); + ++smem_pipe_write_k; + } + + // load Q tile + if (warp_idx_in_warpgroup == 0) { + cutlass::arch::NamedBarrier::sync(NUM_MMA_THREADS + cutlass::NumThreadsPerWarp, + static_cast(NamedBarriers::kQueryEmpty)); + + int lane_predicate = cute::elect_one_sync(); + if (lane_predicate) { + shared_storage.barrier_Q.arrive_and_expect_tx(TmaTransactionBytesQ); + copy(mainloop_params.tma_load_Q.with( + reinterpret_cast( + shared_storage.barrier_Q), + /*mcast_mask=*/0), + tQgQ, tQsQ); + } + } + + shared_storage.barrier_O.wait((work_idx + 1) % 2); + + if (kv_tile_idx == swa_begin_kv_tile_idx) { + pipeline_v.producer_acquire(smem_pipe_write_v); + Tensor tVgViGroup = flatten_1(tVgV(_, _, _, kv_tile_idx)); // (CPY, (CPY_KV, CPY_D)) + Tensor tVsViGroup = + flatten_1(tVsV(_, _, _, smem_pipe_write_v.index())); // (CPY, (CPY_KV, CPY_D)) + copy_if(gmem_tiled_copy_v, v_predicate_fn, tVgViGroup, tVsViGroup); + + pipeline_v.producer_commit(smem_pipe_write_v, cutlass::arch::cpasync_barrier_arrive); + ++smem_pipe_write_v; + } else { + // load second last k-tile and last v-tile + pipeline_k.producer_acquire(smem_pipe_write_k); + Tensor tKgKi = tKgK(_, _, _, kv_tile_idx_decrement(kv_tile_idx)); // (CPY, CPY_KV, CPY_D) + Tensor tKsKi = tKsK(_, _, _, smem_pipe_write_k.index()); // (CPY, CPY_KV, CPY_D) + copy(gmem_tiled_copy_k, tKgKi, tKsKi); + + pipeline_k.producer_commit(smem_pipe_write_k, cutlass::arch::cpasync_barrier_arrive); + ++smem_pipe_write_k; + + pipeline_v.producer_acquire(smem_pipe_write_v); + Tensor tVgViGroup = flatten_1(tVgV(_, _, _, kv_tile_idx)); // (CPY, (CPY_KV, CPY_D)) + Tensor tVsViGroup = + flatten_1(tVsV(_, _, _, smem_pipe_write_v.index())); // (CPY, (CPY_KV, CPY_D)) + copy_if(gmem_tiled_copy_v, v_predicate_fn, tVgViGroup, tVsViGroup); + + pipeline_v.producer_commit(smem_pipe_write_v, cutlass::arch::cpasync_barrier_arrive); + kv_tile_idx = kv_tile_idx_decrement(kv_tile_idx); + ++smem_pipe_write_v; + + // load remaining k/v tiles +#pragma unroll 2 + for (; kv_tile_idx > swa_begin_kv_tile_idx; + kv_tile_idx = kv_tile_idx_decrement(kv_tile_idx)) { + pipeline_k.producer_acquire(smem_pipe_write_k); + + Tensor tKgKi = tKgK(_, _, _, kv_tile_idx_decrement(kv_tile_idx)); // (CPY, CPY_KV, CPY_D) + Tensor tKsKi = tKsK(_, _, _, smem_pipe_write_k.index()); // (CPY, CPY_KV, CPY_D) + copy(gmem_tiled_copy_k, tKgKi, tKsKi); + + pipeline_k.producer_commit(smem_pipe_write_k, cutlass::arch::cpasync_barrier_arrive); + ++smem_pipe_write_k; + + pipeline_v.producer_acquire(smem_pipe_write_v); + Tensor tVgVi = tVgV(_, _, _, kv_tile_idx); // (CPY, CPY_KV, CPY_D) + Tensor tVsVi = tVsV(_, _, _, smem_pipe_write_v.index()); // (CPY, CPY_KV, CPY_D) + copy(gmem_tiled_copy_v, tVgVi, tVsVi); + + pipeline_v.producer_commit(smem_pipe_write_v, cutlass::arch::cpasync_barrier_arrive); + ++smem_pipe_write_v; + } + scheduler.prefetch_next_work(scheduler_params, work_tile_info); + + // load first v tile + { + pipeline_v.producer_acquire(smem_pipe_write_v); + Tensor tVgVi = tVgV(_, _, _, 0); // (CPY, (CPY_KV, CPY_D)) + Tensor tVsVi = tVsV(_, _, _, smem_pipe_write_v.index()); // (CPY, (CPY_KV, CPY_D)) + copy(gmem_tiled_copy_v, tVgVi, tVsVi); + pipeline_v.producer_commit(smem_pipe_write_v, cutlass::arch::cpasync_barrier_arrive); + ++smem_pipe_write_v; + } + } + + scheduler.broadcast_next_work(work_tile_info); + } + + CUTLASS_DEVICE void load_tail(MainloopPipeline pipeline_k, MainloopPipeline pipeline_v, + PipelineState& smem_pipe_write_k, + PipelineState& smem_pipe_write_v) { + pipeline_k.producer_tail(smem_pipe_write_k); + pipeline_v.producer_tail(smem_pipe_write_v); + } +}; + +} // namespace flashinfer + +#endif // FLASHINFER_ATTENTION_HOPPER_SPARSE_MAINLOOP_CUH_ diff --git a/include/flashinfer/attention/hopper/tile_scheduler.cuh b/include/flashinfer/attention/hopper/tile_scheduler.cuh new file mode 100644 index 0000000000000000000000000000000000000000..51a322346f2a7def53f6d09f9bac70604a512905 --- /dev/null +++ b/include/flashinfer/attention/hopper/tile_scheduler.cuh @@ -0,0 +1,322 @@ +/* + * Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri + * Dao. Licensed under the BSD 3-Clause. + * + * Modified by the FlashInfer team. + */ +#ifndef FLASHINFER_ATTENTION_HOPPER_TILE_SCHEDULER_CUH_ +#define FLASHINFER_ATTENTION_HOPPER_TILE_SCHEDULER_CUH_ + +#include "cutlass/arch/barrier.h" +#include "cutlass/fast_math.h" +#include "named_barrier.cuh" + +namespace flashinfer { + +struct SingleTileScheduler { + public: + // Host side kernel arguments + struct Arguments { + int const num_qo_tiles, num_qo_heads, qo_len, kv_len; + cutlass::FastDivmod group_size_fastdiv; + }; + + // Device side kernel params + struct Params { + int const qo_len, kv_len; + cutlass::FastDivmod group_size_fastdiv; + }; + + static Params to_underlying_arguments(Arguments const& args) { + return {args.qo_len, args.kv_len, args.group_size_fastdiv}; + } + + static dim3 get_grid_dim(Arguments const& args, int num_sm) { + return {uint32_t(args.num_qo_tiles), uint32_t(args.num_qo_heads)}; + } + + struct WorkTileInfo { + int q_tile_idx = 0; + int qo_head_idx = 0; + int kv_head_idx = 0; + bool is_valid_tile = false; + + CUTLASS_DEVICE + bool is_valid(Params const& params) const { return is_valid_tile; } + + CUTLASS_DEVICE + auto get_block_coord(Params const& params) const { + return cute::tuple{q_tile_idx, qo_head_idx, kv_head_idx, /*qo_indptr=*/0, + /*kv_indptr=*/0, params.qo_len, params.kv_len, /*batch_idx=*/0}; + } + }; + + CUTLASS_DEVICE + SingleTileScheduler() {} + + CUTLASS_DEVICE + WorkTileInfo get_initial_work(Params const& params) const { + int qo_head_idx = blockIdx.y; + int kv_head_idx = params.group_size_fastdiv.divide(qo_head_idx); + return {/*q_tile_idx=*/int(blockIdx.x), qo_head_idx, kv_head_idx, /*is_valid_tile*/ true}; + } + + CUTLASS_DEVICE + void init_consumer() const {} + + CUTLASS_DEVICE + void prefetch_next_work(Params const& params, WorkTileInfo& current_work) const {} + + CUTLASS_DEVICE + void broadcast_next_work(WorkTileInfo& current_work) const {} + + template + CUTLASS_DEVICE WorkTileInfo get_next_work(Params const& params, + WorkTileInfo const& current_work) const { + return {-1, -1, false}; + } +}; + +template +struct BatchPrefillPersistentTileScheduler { + public: + // Host side kernel arguments + struct Arguments { + IdType *work_indptr, *head_indices, *qo_tile_indices, *qo_indptr, *kv_indptr, *qo_lens, + *kv_lens, *batch_indices; + cutlass::FastDivmod group_size_fastdiv; + int num_qo_heads; // placeholder + }; + + // Device side kernel params + struct Params { + IdType *work_indptr, *head_indices, *qo_tile_indices, *qo_indptr, *kv_indptr, *qo_lens, + *kv_lens, *batch_indices; + cutlass::FastDivmod group_size_fastdiv; + }; + + static Params to_underlying_arguments(Arguments const& args) { + return {args.work_indptr, args.head_indices, args.qo_tile_indices, + args.qo_indptr, args.kv_indptr, args.qo_lens, + args.kv_lens, args.batch_indices, args.group_size_fastdiv}; + } + + static dim3 get_grid_dim(Arguments const& args, int num_sm) { return {(unsigned)num_sm}; } + + struct WorkTileInfo { + int q_tile_idx = 0; + int qo_head_idx = 0; + int kv_head_idx = 0; + int qo_indptr = 0; + int kv_indptr = 0; + int qo_len = 0; + int kv_len = 0; + int batch_idx = 0; + int counter = 0; + int ptr_begin = 0; + int ptr_end = 0; + + CUTLASS_DEVICE + bool is_valid(Params const& params) const { return counter + ptr_begin < ptr_end; } + + CUTLASS_DEVICE + auto get_block_coord(Params const& params) const { + return cute::tuple{q_tile_idx, qo_head_idx, kv_head_idx, qo_indptr, + kv_indptr, qo_len, kv_len, batch_idx}; + } + }; + + CUTLASS_DEVICE + BatchPrefillPersistentTileScheduler() {} + + CUTLASS_DEVICE + WorkTileInfo get_initial_work(Params const& params) const { + int ptr_begin = params.work_indptr[blockIdx.x]; + int ptr_end = params.work_indptr[blockIdx.x + 1]; + if (ptr_begin < ptr_end) { + int work_idx = ptr_begin; + int qo_head_idx = params.head_indices[work_idx]; + int kv_head_idx = params.group_size_fastdiv.divide(qo_head_idx); + return {params.qo_tile_indices[work_idx], + qo_head_idx, + kv_head_idx, + params.qo_indptr[work_idx], + params.kv_indptr[work_idx], + params.qo_lens[work_idx], + params.kv_lens[work_idx], + params.batch_indices[work_idx], + /*counter=*/0, + ptr_begin, + ptr_end}; + } else { + return {-1, -1, -1, -1, -1, -1, -1, 0, ptr_begin, ptr_end}; + } + } + + CUTLASS_DEVICE + void init_consumer() const {} + + CUTLASS_DEVICE + void prefetch_next_work(Params const& params, WorkTileInfo& current_work) const {} + + CUTLASS_DEVICE + void broadcast_next_work(WorkTileInfo& current_work) const {} + + template + CUTLASS_DEVICE WorkTileInfo get_next_work(Params const& params, + WorkTileInfo const& current_work) const { + int work_idx = current_work.ptr_begin + current_work.counter + 1; + if (work_idx < current_work.ptr_end) { + int qo_head_idx = params.head_indices[work_idx]; + int kv_head_idx = params.group_size_fastdiv.divide(qo_head_idx); + return {params.qo_tile_indices[work_idx], + qo_head_idx, + kv_head_idx, + params.qo_indptr[work_idx], + params.kv_indptr[work_idx], + params.qo_lens[work_idx], + params.kv_lens[work_idx], + params.batch_indices[work_idx], + current_work.counter + 1, + current_work.ptr_begin, + current_work.ptr_end}; + } else { + return {-1, + -1, + -1, + -1, + -1, + -1, + -1, + current_work.counter + 1, + current_work.ptr_begin, + current_work.ptr_end}; + } + } +}; + +/*! + * \brief Tile scheduler that maps q/o head to blockIdx.y + */ +template +struct BatchPrefillTileScheduler { + public: + // Host side kernel arguments + struct Arguments { + IdType *work_indptr, *head_indices, *qo_tile_indices, *qo_indptr, *kv_indptr, *qo_lens, + *kv_lens, *batch_indices; // head_indices is a placeholder + cutlass::FastDivmod group_size_fastdiv; + int num_qo_heads; + }; + + // Device side kernel params + struct Params { + IdType *work_indptr, *qo_tile_indices, *qo_indptr, *kv_indptr, *qo_lens, *kv_lens, + *batch_indices; + cutlass::FastDivmod group_size_fastdiv; + int num_qo_heads; + }; + + static Params to_underlying_arguments(Arguments const& args) { + return {args.work_indptr, args.qo_tile_indices, args.qo_indptr, args.kv_indptr, + args.qo_lens, args.kv_lens, args.batch_indices, args.group_size_fastdiv, + args.num_qo_heads}; + } + + static dim3 get_grid_dim(Arguments const& args, int num_sm) { + return {(unsigned)num_sm, (unsigned)args.num_qo_heads}; + } + + struct WorkTileInfo { + int q_tile_idx = 0; + int qo_head_idx = 0; + int kv_head_idx = 0; + int qo_indptr = 0; + int kv_indptr = 0; + int qo_len = 0; + int kv_len = 0; + int batch_idx = 0; + int counter = 0; + int ptr_begin = 0; + int ptr_end = 0; + + CUTLASS_DEVICE + bool is_valid(Params const& params) const { return counter + ptr_begin < ptr_end; } + + CUTLASS_DEVICE + auto get_block_coord(Params const& params) const { + return cute::tuple{q_tile_idx, qo_head_idx, kv_head_idx, qo_indptr, + kv_indptr, qo_len, kv_len, batch_idx}; + } + }; + + CUTLASS_DEVICE + BatchPrefillTileScheduler() {} + + CUTLASS_DEVICE + WorkTileInfo get_initial_work(Params const& params) const { + int ptr_begin = params.work_indptr[blockIdx.x]; + int ptr_end = params.work_indptr[blockIdx.x + 1]; + if (ptr_begin < ptr_end) { + int work_idx = ptr_begin; + int qo_head_idx = blockIdx.y; + int kv_head_idx = params.group_size_fastdiv.divide(qo_head_idx); + return {params.qo_tile_indices[work_idx], + /*qo_head_idx=*/qo_head_idx, + /*kv_head_idx=*/kv_head_idx, + params.qo_indptr[work_idx], + params.kv_indptr[work_idx], + params.qo_lens[work_idx], + params.kv_lens[work_idx], + params.batch_indices[work_idx], + /*counter=*/0, + ptr_begin, + ptr_end}; + } else { + return {-1, -1, -1, -1, -1, -1, -1, 0, ptr_begin, ptr_end}; + } + } + + CUTLASS_DEVICE + void init_consumer() const {} + + CUTLASS_DEVICE + void prefetch_next_work(Params const& params, WorkTileInfo& current_work) const {} + + CUTLASS_DEVICE + void broadcast_next_work(WorkTileInfo& current_work) const {} + + template + CUTLASS_DEVICE WorkTileInfo get_next_work(Params const& params, + WorkTileInfo const& current_work) const { + int work_idx = current_work.ptr_begin + current_work.counter + 1; + if (work_idx < current_work.ptr_end) { + return {params.qo_tile_indices[work_idx], + current_work.qo_head_idx, + current_work.kv_head_idx, + params.qo_indptr[work_idx], + params.kv_indptr[work_idx], + params.qo_lens[work_idx], + params.kv_lens[work_idx], + params.batch_indices[work_idx], + current_work.counter + 1, + current_work.ptr_begin, + current_work.ptr_end}; + } else { + return {-1, + -1, + -1, + -1, + -1, + -1, + -1, + current_work.counter + 1, + current_work.ptr_begin, + current_work.ptr_end}; + } + } +}; + +} // namespace flashinfer + +#endif // FLASHINFER_ATTENTION_HOPPER_TILE_SCHEDULER_CUH_ diff --git a/include/flashinfer/attention/hopper/utils.cuh b/include/flashinfer/attention/hopper/utils.cuh new file mode 100644 index 0000000000000000000000000000000000000000..8aeb5b194426a770c379df90801dea69bb528a81 --- /dev/null +++ b/include/flashinfer/attention/hopper/utils.cuh @@ -0,0 +1,197 @@ +/* + * Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri + * Dao. Licensed under the BSD 3-Clause. + * + * Modified by the FlashInfer team. + */ +#ifndef FLASHINFER_ATTENTION_HOPPER_UTILS_CUH_ +#define FLASHINFER_ATTENTION_HOPPER_UTILS_CUH_ + +#include +#include +#include +#include + +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 +#include +#endif + +#include +#include +#include +#include +#include + +#include +#include +#include + +#include "../../math.cuh" +#include "../../utils.cuh" +#include "cutlass/fast_math.h" + +namespace flashinfer { + +using namespace cute; + +template +CUTLASS_DEVICE int get_swa_begin_kv_tile_idx(int window_left, int q_tile_idx, const int qo_len, + const int kv_len) { + return std::max((q_tile_idx * CTA_Q + kv_len - qo_len - window_left) / CTA_KV - 1, 0); +} + +template +CUTLASS_DEVICE int get_swa_end_kv_tile_idx(int window_left, int q_tile_idx, const int qo_len, + const int kv_len) { + return std::max(((q_tile_idx + 1) * CTA_Q + kv_len - qo_len - window_left) / CTA_KV, -1); +} + +template +CUTLASS_HOST_DEVICE auto flatten_1(TensorT tensor) { + Tensor tensor_flatten = cute::flatten(tensor); + return cute::group_modes<1, rank(tensor_flatten)>(tensor_flatten); +} + +CUTLASS_HOST_DEVICE auto get_gmem_layout(int nnz, int num_heads, int head_dim, int64_t n_stride, + int64_t h_stride) { + return make_layout(make_shape(nnz, head_dim, num_heads), + make_stride(n_stride, cute::_1{}, h_stride)); +} + +CUTLASS_HOST_DEVICE auto get_lse_gmem_layout(int nnz, int num_heads) { + return make_layout(make_shape(num_heads, nnz), make_stride(cute::_1{}, int64_t(num_heads))); +} + +template +CUTLASS_DEVICE auto get_local_tile_tensor(const MTensor& m_tensor, const Shape& tile_shape, + int head_idx, int offset, int seq_len) { + auto g_offset = local_tile(m_tensor(_, _, head_idx), cute::make_shape(1, get<1>(tile_shape)), + make_coord(offset, _0{})); + auto g_sequence = + make_tensor(g_offset.data(), + make_layout(cute::make_shape(seq_len, get<1>(tile_shape)), g_offset.stride())); + auto g_tensor = local_tile(g_sequence, tile_shape, make_coord(_, _0{})); + return g_tensor; +} + +template +CUTLASS_DEVICE auto get_lse_local_tile_tensor(const MTensor& m_tensor, const Shape& tile_shape, + int head_idx, int offset, int seq_len) { + auto g_offset = local_tile(m_tensor(head_idx, _), cute::make_shape(_1{}), make_coord(offset)); + + auto g_sequence = make_tensor(g_offset.data(), make_layout(cute::make_shape(seq_len), + cute::make_shape(shape<0>(m_tensor)))); + auto g_tensor = local_tile(g_sequence, tile_shape, make_coord(_)); + return g_tensor; +} + +// For SM90, convert acc_layout from ((2, 2, V), MMA_M, MMA_N) to (nrow=(2, MMA_M), ncol=(2, V, +// MMA_N)) +template +__forceinline__ __device__ auto convert_layout_acc_rowcol(Layout acc_layout) { + static_assert(decltype(size<0, 0>(acc_layout))::value == 2); + static_assert(decltype(size<0, 1>(acc_layout))::value == 2); + static_assert(decltype(rank(acc_layout))::value == 3); + auto l = acc_layout; + return make_layout(make_layout(get<0, 1>(l), get<1>(l)), + make_layout(get<0, 0>(l), get<0, 2>(l), get<2>(l))); +}; + +// For SM90, convert acc_layout from ((2, 2, N / 8), MMA_M, MMA_N) to ((2, 2, 2), MMA_M, (N / 16, +// MMA_N)) +template +__forceinline__ __device__ auto convert_layout_acc_Aregs(Layout acc_layout) { + using X = Underscore; + static_assert(decltype(size<0, 0>(acc_layout))::value == 2); + static_assert(decltype(size<0, 1>(acc_layout))::value == 2); + static_assert(decltype(rank(acc_layout))::value == 3); + static_assert(decltype(rank(get<0>(acc_layout)))::value == 3); + auto l = logical_divide(get<0>(acc_layout), Shape{}); // (2, 2, (2, N / 16))) + return make_layout(make_layout(get<0>(l), get<1>(l), get<2, 0>(l)), get<1>(acc_layout), + make_layout(get<2, 1>(l), get<2>(acc_layout))); +}; + +// Convert acc_layout from ((2, 2, N / 8), MMA_M, MMA_N) to ((4, 2, 2), MMA_M, +// (N / 32, MMA_N)) +template +__forceinline__ __device__ auto convert_layout_acc_Aregs_fp8(Layout acc_layout) { + using X = Underscore; + static_assert(decltype(size<0, 0>(acc_layout))::value == 2); + static_assert(decltype(size<0, 1>(acc_layout))::value == 2); + static_assert(decltype(rank(acc_layout))::value == 3); + static_assert(decltype(rank(get<0>(acc_layout)))::value == 3); + auto l = logical_divide(get<0>(acc_layout), Shape{}); // (2, 2, (2, N / 32))) + return make_layout(make_layout(Shape<_4, _2, _2>{}), get<1>(acc_layout), + make_layout(get<2, 1>(l), get<2>(acc_layout))); +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// Byte permute for fp8 kernel +template +CUTLASS_DEVICE void permute_regs_A_to_C(Fragment& accum) { + auto data = accum.data(); +#pragma unroll + for (int n = 0; n < size(accum); n += 8) { + uint32_t* data_32bit = reinterpret_cast(&data[n]); + auto upper = data_32bit[0]; + auto lower = data_32bit[1]; + data_32bit[0] = __byte_perm(upper, lower, 0x5410); + data_32bit[1] = __byte_perm(upper, lower, 0x7632); + } +} + +template +__forceinline__ __device__ auto convert_type(Tensor const& tensor) { + using From_type = typename Engine::value_type; + constexpr int numel = decltype(size(tensor))::value; + cutlass::NumericArrayConverter + convert_op; + // HACK: this requires tensor to be "contiguous" + auto frag = convert_op(*reinterpret_cast*>(tensor.data())); + return make_tensor(make_rmem_ptr(&frag), tensor.layout()); +} + +template +__forceinline__ __device__ void gemm(TiledMma& tiled_mma, TensorA const& tCrA, TensorB const& tCrB, + TensorC& tCrC) { + constexpr bool Is_RS = + !cute::is_base_of::value; + // Need to cast away const on tCrA since warpgroup_fence_operand doesn't take const + if constexpr (Is_RS) { + warpgroup_fence_operand(const_cast(tCrA)); + } + warpgroup_fence_operand(tCrC); + warpgroup_arrive(); + if constexpr (init) { + tiled_mma.accumulate_ = GMMA::ScaleOut::Zero; + // Unroll the K mode manually to set scale D to 1 + CUTLASS_PRAGMA_UNROLL + for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) { + cute::gemm(tiled_mma, tCrA(_, _, k_block), tCrB(_, _, k_block), tCrC); + tiled_mma.accumulate_ = GMMA::ScaleOut::One; + } + } else { + // cute::gemm(tiled_mma, tCrA, tCrB, tCrC); + // Unroll the K mode manually to set scale D to 1 + CUTLASS_PRAGMA_UNROLL + for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) { + cute::gemm(tiled_mma, tCrA(_, _, k_block), tCrB(_, _, k_block), tCrC); + tiled_mma.accumulate_ = GMMA::ScaleOut::One; + } + } + warpgroup_commit_batch(); + if constexpr (wg_wait >= 0) { + warpgroup_wait(); + } + warpgroup_fence_operand(tCrC); + if constexpr (Is_RS) { + warpgroup_fence_operand(const_cast(tCrA)); + } +} + +} // namespace flashinfer + +#endif // FLASHINFER_ATTENTION_HOPPER_UTILS_CUH_ diff --git a/include/flashinfer/attention/hopper/variant_helper.cuh b/include/flashinfer/attention/hopper/variant_helper.cuh new file mode 100644 index 0000000000000000000000000000000000000000..22fda01f4a07924ff1d83db7c6a7372381a1dac7 --- /dev/null +++ b/include/flashinfer/attention/hopper/variant_helper.cuh @@ -0,0 +1,64 @@ +/* + * Copyright (c) 2024 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef FLASHINFER_ATTENTION_HOPPER_VARIANT_HELPER_H +#define FLASHINFER_ATTENTION_HOPPER_VARIANT_HELPER_H + +#include + +#include + +namespace flashinfer { + +#define REGISTER_QUERY_TRANSFORM(params, q, ...) \ + template \ + __device__ __forceinline__ T QueryTransform(const MainloopParams& params, void* q_smem) { \ + __VA_ARGS__ \ + } + +#define REGISTER_KEY_TRANSFORM(params, k, ...) \ + template \ + __device__ __forceinline__ T KeyTransform(const MainloopParams& params, void* k_smem) { \ + __VA_ARGS__ \ + } + +#define REGISTER_LOGITS_TRANSFORM(params, logits, batch_idx, qo_idx, kv_idx, qo_head_idx, \ + kv_head_idx, ...) \ + template \ + __device__ __forceinline__ T LogitsTransform( \ + const MainloopParams& params, T logits, uint32_t batch_idx, uint32_t qo_idx, \ + uint32_t kv_idx, uint32_t qo_head_idx, uint32_t kv_head_idx) { \ + __VA_ARGS__ \ + } + +#define REGISTER_LOGITS_MASK(params, batch_idx, qo_idx, kv_idx, qo_head_idx, kv_head_idx, ...) \ + template \ + __device__ __forceinline__ bool LogitsMask(const MainloopParams& params, uint32_t batch_idx, \ + uint32_t qo_idx, uint32_t kv_idx, \ + uint32_t qo_head_idx, uint32_t kv_head_idx) { \ + __VA_ARGS__ \ + } + +struct AttentionVariantBase { + REGISTER_LOGITS_TRANSFORM(params, logits, batch_idx, qo_idx, kv_idx, qo_head_idx, kv_head_idx, + { return logits; }) + + REGISTER_LOGITS_MASK(params, batch_idx, qo_idx, kv_idx, qo_head_idx, kv_head_idx, + { return true; }) +}; + +} // namespace flashinfer + +#endif // FLASHINFER_ATTENTION_HOPPER_VARIANT_HELPER_H diff --git a/include/flashinfer/attention/hopper/variants.cuh b/include/flashinfer/attention/hopper/variants.cuh new file mode 100644 index 0000000000000000000000000000000000000000..1a199e2d3054046c4c95e087142cec23960a04b2 --- /dev/null +++ b/include/flashinfer/attention/hopper/variants.cuh @@ -0,0 +1,108 @@ +/* + * Copyright (c) 2024 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +// NOTE(Zihao): we should merge this with include/flashinfer/attention/variants.cuh in the future +#ifndef FLASHINFER_ATTENTION_HOPPER_VARIANTS_CUH_ +#define FLASHINFER_ATTENTION_HOPPER_VARIANTS_CUH_ +#include + +#include "../../math.cuh" +#include "attention_updater.cuh" +#include "variant_helper.cuh" + +namespace flashinfer { + +struct StandardAttention { + float sm_scale_log2; + + template + __device__ StandardAttention(const MainloopParams& params, const BlockCoord& block_coord) { + sm_scale_log2 = params.additional_params.sm_scale * math::log2e; + } + + template + __device__ auto GetAttentionUpdater() { + return OnlineSoftmax(sm_scale_log2); + } + + REGISTER_LOGITS_TRANSFORM(params, logits, batch_idx, qo_idx, kv_idx, qo_head_idx, kv_head_idx, + { return logits; }) +}; + +struct LogitsSoftCap { + float pre_tanh_scale; + float post_tanh_scale; + + template + __device__ LogitsSoftCap(const MainloopParams& params, const BlockCoord& block_coord) { + pre_tanh_scale = + params.additional_params.sm_scale * math::ptx_rcp(params.additional_params.logits_soft_cap); + post_tanh_scale = math::log2e * params.additional_params.logits_soft_cap; + } + + template + __device__ auto GetAttentionUpdater() { + return OnlineSoftmax(post_tanh_scale); + } + + REGISTER_LOGITS_TRANSFORM(params, logits, batch_idx, qo_idx, kv_idx, qo_head_idx, kv_head_idx, + { return math::tanh(logits * pre_tanh_scale); }) +}; + +struct StandardFP8Attention { + float p_scale, scale_pv, sm_scale_with_qk_log2; + + template + __device__ StandardFP8Attention(const MainloopParams& params, const BlockCoord& block_coord) { + auto [q_tile_idx, qo_head_idx, kv_head_idx, qo_indptr, kv_indptr, qo_len, kv_len, batch_idx] = + block_coord; + // 448 for e4m3; 57344 for e5m2 + p_scale = std::numeric_limits::max(); + scale_pv = params.additional_params.scale_v[kv_head_idx] / p_scale; + sm_scale_with_qk_log2 = params.additional_params.scale_q[qo_head_idx] * + params.additional_params.scale_k[kv_head_idx] * + params.additional_params.sm_scale * math::log2e; + } + + template + __device__ auto GetAttentionUpdater() { + return OnlineSoftmax(sm_scale_with_qk_log2); + } + + template + __device__ __forceinline__ void PQuantize(Tensor0& tSrS) { +#pragma unroll + for (int i = 0; i < size(tSrS); ++i) { + tSrS(i) *= p_scale; + } + } + + template + __device__ __forceinline__ void ODequantize(const MainloopParams& params, Tensor0& tOrO, + uint32_t qo_head_idx, uint32_t kv_head_idx) { + // we fuse the PV dequantization into online_softmax.finalize + } + + REGISTER_LOGITS_TRANSFORM(params, logits, batch_idx, qo_idx, kv_idx, qo_head_idx, kv_head_idx, + { return logits; }) +}; + +template +using DefaultAttention = std::conditional_t; +using DefaultFP8Attention = StandardFP8Attention; + +} // namespace flashinfer + +#endif // FLASHINFER_ATTENTION_HOPPER_VARIANTS_CUH_ diff --git a/include/flashinfer/attention/mask.cuh b/include/flashinfer/attention/mask.cuh new file mode 100644 index 0000000000000000000000000000000000000000..6692b0cf3f54c708f134055f49d7d814de376b58 --- /dev/null +++ b/include/flashinfer/attention/mask.cuh @@ -0,0 +1,30 @@ +/* + * Copyright (c) 2024 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef FLASHINFER_ATTENTION_MASK_CUH_ +#define FLASHINFER_ATTENTION_MASK_CUH_ + +namespace flashinfer { + +enum class MaskMode { + kNone = 0U, // No mask + kCausal = 1U, // Causal mask + kCustom = 2U, // Custom mask + kMultiItemScoring = 3U, +}; + +} // namespace flashinfer + +#endif // FLASHINFER_ATTENTION_MASK_CUH_ diff --git a/include/flashinfer/attention/mla.cuh b/include/flashinfer/attention/mla.cuh new file mode 100644 index 0000000000000000000000000000000000000000..31401ff1c57a470b4f4926cc7d5bb7f48dc870dc --- /dev/null +++ b/include/flashinfer/attention/mla.cuh @@ -0,0 +1,1048 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef FLASHINFER_MLA_FA2_CUH_ +#define FLASHINFER_MLA_FA2_CUH_ +#include + +#include +#include + +#include "../profiler.cuh" +#include "mla_params.cuh" +#include "prefill.cuh" +#include "variant_helper.cuh" + +namespace flashinfer { + +namespace mla { + +struct StandardAttention : AttentionVariantBase { + float sm_scale_log2; + + PROFILER_CLOSURE_PARAMS_DECL + + template + __device__ __host__ StandardAttention(const Params& params, uint32_t batch_idx, + uint8_t* smem_ptr) { + sm_scale_log2 = params.sm_scale * math::log2e; + } +}; + +template +struct SharedStorageQKVO { + union { + struct { + alignas(16) DTypeQ q_smem_nope[CTA_TILE_Q * HEAD_DIM_CKV]; + alignas(16) DTypeQ q_smem_pe[CTA_TILE_Q * HEAD_DIM_KPE]; + alignas(16) DTypeKV ckv_smem[NUM_STAGES][CTA_TILE_KV * HEAD_DIM_CKV]; + alignas(16) DTypeKV + kpe_p_smem[NUM_STAGES] + [CTA_TILE_KV * (HEAD_DIM_KPE > CTA_TILE_Q ? HEAD_DIM_KPE : CTA_TILE_Q)]; + union { + alignas(16) float m_wg[2][CTA_TILE_Q]; // cross warpgroup synchronization + alignas(16) float d_wg[2][CTA_TILE_Q]; // cross warpgroup synchronization + }; + }; + alignas(16) DTypeO o_smem[CTA_TILE_Q * HEAD_DIM_CKV]; + }; +}; + +template +struct KernelTraits { + static constexpr bool CAUSAL = CAUSAL_; + static constexpr uint32_t NUM_STAGES = NUM_STAGES_; + // NOTE(Zihao): whether to shard Q*K computation across warpgroups + // if true, each warpgroup will compute a subset of Q*K (sharded on the KV dimension) + // if false, each warpgroup will compute the full Q*K, which is duplicated across warpgroups + static constexpr bool QK_SHARD = QK_SHARD_; + static constexpr uint32_t NUM_MMA_KV = CTA_TILE_KV_ / 16; + static constexpr uint32_t HEAD_DIM_CKV = HEAD_DIM_CKV_; + static constexpr uint32_t HEAD_DIM_KPE = HEAD_DIM_KPE_; + static constexpr uint32_t HEAD_DIM_ALL = HEAD_DIM_CKV + HEAD_DIM_KPE; + static constexpr uint32_t NUM_MMA_D_CKV = HEAD_DIM_CKV / 16; + static constexpr uint32_t NUM_MMA_D_KPE = HEAD_DIM_KPE / 16; + static constexpr uint32_t NUM_THREADS = 256; + static constexpr uint32_t CTA_TILE_Q = CTA_TILE_Q_; + static constexpr uint32_t CTA_TILE_KV = CTA_TILE_KV_; + + static constexpr SwizzleMode SWIZZLE_MODE_Q_NOPE = SwizzleMode::k128B; + static constexpr SwizzleMode SWIZZLE_MODE_Q_PE = SwizzleMode::k128B; + static constexpr SwizzleMode SWIZZLE_MODE_CKV = SwizzleMode::k128B; + static constexpr SwizzleMode SWIZZLE_MODE_KPE = SwizzleMode::k128B; + static constexpr SwizzleMode SWIZZLE_MODE_P = + CTA_TILE_KV >= 64 ? SwizzleMode::k128B : SwizzleMode::k64B; + static constexpr SwizzleMode SWIZZLE_MODE_O = SwizzleMode::k128B; + static constexpr uint32_t UPCAST_STRIDE_Q_NOPE = HEAD_DIM_CKV / upcast_size(); + static constexpr uint32_t UPCAST_STRIDE_Q_PE = HEAD_DIM_KPE / upcast_size(); + static constexpr uint32_t UPCAST_STRIDE_CKV = HEAD_DIM_CKV / upcast_size(); + static constexpr uint32_t UPCAST_STRIDE_KPE = HEAD_DIM_KPE / upcast_size(); + static constexpr uint32_t UPCAST_STRIDE_FINAL_O = HEAD_DIM_CKV / upcast_size(); + static constexpr uint32_t UPCAST_STRIDE_P = CTA_TILE_KV / upcast_size(); + + using DTypeQ = DTypeQ_; + using DTypeKV = DTypeKV_; + using DTypeO = DTypeO_; + using IdType = IdType_; + using DTypeQKAccum = float; + + using SharedStorage = SharedStorageQKVO; + using AttentionVariant = StandardAttention; + + static constexpr DTypeQKAccum MaskFillValue = -math::inf; +}; + +template +__device__ __forceinline__ void init_states_(float (*o_frag)[8], typename KTraits::DTypeQKAccum* m, + float* d) { +#pragma unroll + for (uint32_t mma_d = 0; mma_d < KTraits::NUM_MMA_D_CKV / 2; ++mma_d) { +#pragma unroll + for (uint32_t reg_id = 0; reg_id < 8; ++reg_id) { + o_frag[mma_d][reg_id] = 0.f; + } + } + +#pragma unroll + for (uint32_t j = 0; j < 2; ++j) { + m[j] = typename KTraits::DTypeQKAccum(-math::inf); + d[j] = 1.f; + } +} + +template +__device__ __forceinline__ void load_q( + typename KTraits::SharedStorage* smem_storage, typename KTraits::DTypeQ* q_nope, + typename KTraits::DTypeQ* q_pe, const uint32_t q_nope_stride_n, const uint32_t q_nope_stride_h, + const uint32_t q_pe_stride_n, const uint32_t q_pe_stride_h, const uint32_t q_len, + const uint32_t packed_offset, const uint_fastdiv& num_heads) { + using DTypeQ = typename KTraits::DTypeQ; + constexpr uint32_t UPCAST_STRIDE_Q_NOPE = KTraits::UPCAST_STRIDE_Q_NOPE; + constexpr uint32_t UPCAST_STRIDE_Q_PE = KTraits::UPCAST_STRIDE_Q_PE; + constexpr uint32_t NUM_MMA_D_CKV = KTraits::NUM_MMA_D_CKV; + constexpr uint32_t NUM_MMA_D_KPE = KTraits::NUM_MMA_D_KPE; + const uint32_t lane_idx = threadIdx.x; + const uint32_t warpgroup_idx = threadIdx.z; + const uint32_t warp_idx_in_wg = threadIdx.y; + + smem_t q_smem_nope(smem_storage->q_smem_nope); + smem_t q_smem_pe(smem_storage->q_smem_pe); + +#pragma unroll + for (uint32_t mma_q = 0; mma_q < 2; ++mma_q) { + uint32_t q, r; + num_heads.divmod( + packed_offset + lane_idx / 8 + (warpgroup_idx + mma_q * 2) * 16 + warp_idx_in_wg * 4, q, r); + DTypeQ* q_nope_ptr = + q_nope + q * q_nope_stride_n + r * q_nope_stride_h + (lane_idx % 8) * upcast_size(); + DTypeQ* q_pe_ptr = + q_pe + q * q_pe_stride_n + r * q_pe_stride_h + (lane_idx % 8) * upcast_size(); +#pragma unroll + for (uint32_t mma_d = 0; mma_d < KTraits::NUM_MMA_D_CKV / 4; ++mma_d) { + uint32_t q_smem_nope_offset_w = + q_smem_nope.template get_permuted_offset( + 32 * mma_q + warpgroup_idx * 16 + warp_idx_in_wg * 4 + lane_idx / 8, + mma_d * 8 + lane_idx % 8); + q_smem_nope.load_128b_async(q_smem_nope_offset_w, q_nope_ptr, + q < q_len); + q_nope_ptr += 8 * upcast_size(); + } +#pragma unroll + for (uint32_t mma_d = 0; mma_d < KTraits::NUM_MMA_D_KPE / 4; ++mma_d) { + uint32_t q_smem_pe_offset_w = q_smem_pe.template get_permuted_offset( + 32 * mma_q + warpgroup_idx * 16 + warp_idx_in_wg * 4 + lane_idx / 8, + mma_d * 8 + lane_idx % 8); + + q_smem_pe.load_128b_async(q_smem_pe_offset_w, q_pe_ptr, + q < q_len); + q_pe_ptr += 8 * upcast_size(); + } + } +} + +template +__device__ __forceinline__ void load_kv( + typename KTraits::SharedStorage* smem_storage, typename KTraits::DTypeKV* ckv, + typename KTraits::DTypeKV* kpe, typename KTraits::IdType* indices, const uint32_t ckv_stride_n, + const uint32_t ckv_stride_page, const uint32_t kpe_stride_n, const uint32_t kpe_stride_page, + const uint32_t packed_kv_bound, const uint32_t packed_block_iter_base, + const uint_fastdiv& block_size, const uint32_t stage_idx) { + using DTypeKV = typename KTraits::DTypeKV; + constexpr uint32_t UPCAST_STRIDE_CKV = KTraits::UPCAST_STRIDE_CKV; + constexpr uint32_t UPCAST_STRIDE_KPE = KTraits::UPCAST_STRIDE_KPE; + constexpr uint32_t NUM_MMA_D_CKV = KTraits::NUM_MMA_D_CKV; + constexpr uint32_t NUM_MMA_D_KPE = KTraits::NUM_MMA_D_KPE; + const uint32_t lane_idx = threadIdx.x; + const uint32_t warpgroup_idx = threadIdx.z; + const uint32_t warp_idx_in_wg = threadIdx.y; + + smem_t ckv_smem(smem_storage->ckv_smem[stage_idx]); + smem_t kpe_smem(smem_storage->kpe_p_smem[stage_idx]); + + if constexpr (KTraits::NUM_MMA_KV == 1) { + if (warpgroup_idx == 0) { + uint32_t q, r; + uint32_t packed_block_iter = packed_block_iter_base + lane_idx / 8 + warp_idx_in_wg * 4; + block_size.divmod(packed_block_iter, q, r); + + DTypeKV* ckv_ptr = ckv + + (packed_block_iter < packed_kv_bound ? indices[q] : 0) * ckv_stride_page + + r * ckv_stride_n + (lane_idx % 8) * upcast_size(); + DTypeKV* kpe_ptr = kpe + + (packed_block_iter < packed_kv_bound ? indices[q] : 0) * kpe_stride_page + + r * kpe_stride_n + (lane_idx % 8) * upcast_size(); + +#pragma unroll + for (uint32_t mma_d = 0; mma_d < KTraits::NUM_MMA_D_CKV / 4; ++mma_d) { + uint32_t ckv_smem_offset_w = ckv_smem.template get_permuted_offset( + warp_idx_in_wg * 4 + lane_idx / 8, 8 * mma_d + lane_idx % 8); + ckv_smem.load_128b_async(ckv_smem_offset_w, ckv_ptr, + packed_block_iter < packed_kv_bound); + ckv_ptr += 8 * upcast_size(); + } + +#pragma unroll + for (uint32_t mma_d = 0; mma_d < KTraits::NUM_MMA_D_KPE / 4; ++mma_d) { + uint32_t kpe_smem_offset_w = kpe_smem.template get_permuted_offset( + warp_idx_in_wg * 4 + lane_idx / 8, 8 * mma_d + lane_idx % 8); + kpe_smem.load_128b_async(kpe_smem_offset_w, kpe_ptr, + packed_block_iter < packed_kv_bound); + kpe_ptr += 8 * upcast_size(); + } + } + } else { +#pragma unroll + for (uint32_t mma_kv = 0; mma_kv < KTraits::NUM_MMA_KV / 2; ++mma_kv) { + uint32_t q, r; + uint32_t packed_block_iter = packed_block_iter_base + lane_idx / 8 + + (warpgroup_idx + mma_kv * 2) * 16 + warp_idx_in_wg * 4; + block_size.divmod(packed_block_iter, q, r); + + DTypeKV* ckv_ptr = ckv + + (packed_block_iter < packed_kv_bound ? indices[q] : 0) * ckv_stride_page + + r * ckv_stride_n + (lane_idx % 8) * upcast_size(); + DTypeKV* kpe_ptr = kpe + + (packed_block_iter < packed_kv_bound ? indices[q] : 0) * kpe_stride_page + + r * kpe_stride_n + (lane_idx % 8) * upcast_size(); + +#pragma unroll + for (uint32_t mma_d = 0; mma_d < KTraits::NUM_MMA_D_CKV / 4; ++mma_d) { + uint32_t ckv_smem_offset_w = ckv_smem.template get_permuted_offset( + 32 * mma_kv + warpgroup_idx * 16 + warp_idx_in_wg * 4 + lane_idx / 8, + 8 * mma_d + lane_idx % 8); + ckv_smem.load_128b_async(ckv_smem_offset_w, ckv_ptr, + packed_block_iter < packed_kv_bound); + ckv_ptr += 8 * upcast_size(); + } + +#pragma unroll + for (uint32_t mma_d = 0; mma_d < KTraits::NUM_MMA_D_KPE / 4; ++mma_d) { + uint32_t kpe_smem_offset_w = kpe_smem.template get_permuted_offset( + 32 * mma_kv + warpgroup_idx * 16 + warp_idx_in_wg * 4 + lane_idx / 8, + 8 * mma_d + lane_idx % 8); + kpe_smem.load_128b_async(kpe_smem_offset_w, kpe_ptr, + packed_block_iter < packed_kv_bound); + kpe_ptr += 8 * upcast_size(); + } + } + } +} + +template +__device__ __forceinline__ void compute_qk_(smem_t q_smem, + smem_t k_smem, + typename KTraits::DTypeQKAccum (*s_frag)[8]) { + const uint32_t lane_idx = threadIdx.x, warpgroup_idx = threadIdx.z, warp_idx_in_wg = threadIdx.y; + alignas(16) uint32_t q_frag[4], k_frag[4]; + // compute q*k^T +#pragma unroll + for (uint32_t mma_d = 0; mma_d < NUM_MMA_D_QK; ++mma_d) { + uint32_t q_smem_offset_r = q_smem.template get_permuted_offset( + warp_idx_in_wg * 16 + lane_idx % 16, mma_d * 2 + lane_idx / 16); + q_smem.ldmatrix_m8n8x4(q_smem_offset_r, q_frag); + + if constexpr (KTraits::QK_SHARD) { +#pragma unroll + for (uint32_t mma_kv = 0; mma_kv < KTraits::NUM_MMA_KV / 2; ++mma_kv) { + uint32_t k_smem_offset_r = k_smem.template get_permuted_offset( + (warpgroup_idx * (KTraits::NUM_MMA_KV / 2) + mma_kv) * 16 + 8 * (lane_idx / 16) + + lane_idx % 8, + 2 * mma_d + (lane_idx % 16) / 8); + + k_smem.ldmatrix_m8n8x4(k_smem_offset_r, k_frag); + + if (init && mma_d == 0) { + mma::mma_sync_m16n16k16_row_col_f16f16f32( + s_frag[mma_kv], q_frag, k_frag); + } else { + mma::mma_sync_m16n16k16_row_col_f16f16f32(s_frag[mma_kv], + q_frag, k_frag); + } + } + } else { +#pragma unroll + for (uint32_t mma_kv = 0; mma_kv < KTraits::NUM_MMA_KV; ++mma_kv) { + uint32_t k_smem_offset_r = k_smem.template get_permuted_offset( + mma_kv * 16 + 8 * (lane_idx / 16) + lane_idx % 8, 2 * mma_d + (lane_idx % 16) / 8); + + k_smem.ldmatrix_m8n8x4(k_smem_offset_r, k_frag); + + if (init && mma_d == 0) { + mma::mma_sync_m16n16k16_row_col_f16f16f32( + s_frag[mma_kv], q_frag, k_frag); + } else { + mma::mma_sync_m16n16k16_row_col_f16f16f32(s_frag[mma_kv], + q_frag, k_frag); + } + } + } + } +} + +template +__device__ __forceinline__ void logits_mask_(const uint32_t qo_packed_idx_base, + const uint32_t kv_idx_base, const uint32_t qo_len, + const uint32_t kv_len, const uint32_t kv_end, + const uint_fastdiv num_heads, + typename KTraits::DTypeQKAccum (*s_frag)[8]) { + const uint32_t lane_idx = threadIdx.x, warpgroup_idx = threadIdx.z, warp_idx_in_wg = threadIdx.y; + constexpr uint32_t NUM_MMA_KV = KTraits::NUM_MMA_KV; + using DTypeQKAccum = typename KTraits::DTypeQKAccum; + uint32_t q[2]; +#pragma unroll + for (uint32_t j = 0; j < 2; ++j) { + q[j] = (qo_packed_idx_base + warp_idx_in_wg * 16 + lane_idx / 4 + 8 * j) / num_heads; + } + + if constexpr (KTraits::QK_SHARD) { +#pragma unroll + for (uint32_t mma_kv = 0; mma_kv < NUM_MMA_KV / 2; ++mma_kv) { +#pragma unroll + for (uint32_t reg_id = 0; reg_id < 8; ++reg_id) { + const uint32_t q_idx = q[(reg_id % 4) / 2], + kv_idx = kv_idx_base + warpgroup_idx * (NUM_MMA_KV / 2) * 16 + mma_kv * 16 + + 2 * (lane_idx % 4) + 8 * (reg_id / 4) + reg_id % 2; + const bool mask = + (!(KTraits::CAUSAL ? (kv_idx + qo_len > kv_len + q_idx || (kv_idx >= kv_end)) + : kv_idx >= kv_end)); + s_frag[mma_kv][reg_id] = (mask) ? s_frag[mma_kv][reg_id] : (KTraits::MaskFillValue); + } + } + } else { +#pragma unroll + for (uint32_t mma_kv = 0; mma_kv < NUM_MMA_KV; ++mma_kv) { +#pragma unroll + for (uint32_t reg_id = 0; reg_id < 8; ++reg_id) { + const uint32_t q_idx = q[(reg_id % 4) / 2], kv_idx = kv_idx_base + mma_kv * 16 + + 2 * (lane_idx % 4) + 8 * (reg_id / 4) + + reg_id % 2; + const bool mask = + (!(KTraits::CAUSAL ? (kv_idx + qo_len > kv_len + q_idx || (kv_idx >= kv_end)) + : kv_idx >= kv_end)); + s_frag[mma_kv][reg_id] = (mask) ? s_frag[mma_kv][reg_id] : (KTraits::MaskFillValue); + } + } + } +} + +template +__device__ __forceinline__ void update_mdo_states_(typename KTraits::SharedStorage* smem_storage, + const uint32_t stage_idx, + typename KTraits::AttentionVariant variant, + typename KTraits::DTypeQKAccum (*s_frag)[8], + float (*o_frag)[8], + typename KTraits::DTypeQKAccum* m, float* d) { + using DTypeQKAccum = typename KTraits::DTypeQKAccum; + using AttentionVariant = typename KTraits::AttentionVariant; + const float sm_scale = variant.sm_scale_log2; + const uint32_t warpgroup_idx = threadIdx.z, lane_idx = threadIdx.x, warp_idx_in_wg = threadIdx.y; + float m_prev[2]; + if constexpr (KTraits::QK_SHARD) { +#pragma unroll + for (uint32_t j = 0; j < 2; ++j) { + m_prev[j] = m[j]; +#pragma unroll + for (uint32_t mma_kv = 0; mma_kv < KTraits::NUM_MMA_KV / 2; ++mma_kv) { + float m_local = max(max(s_frag[mma_kv][j * 2 + 0], s_frag[mma_kv][j * 2 + 1]), + max(s_frag[mma_kv][j * 2 + 4], s_frag[mma_kv][j * 2 + 5])); + m[j] = max(m[j], m_local); + } + m[j] = max(m[j], math::shfl_xor_sync(m[j], 0x2)); + m[j] = max(m[j], math::shfl_xor_sync(m[j], 0x1)); + if (lane_idx % 4 == 0) { + smem_storage->m_wg[warpgroup_idx][warp_idx_in_wg * 16 + j * 8 + lane_idx / 4] = m[j]; + } + } + + __syncthreads(); + +#pragma unroll + for (uint32_t j = 0; j < 2; ++j) { + m[j] = max(smem_storage->m_wg[0][warp_idx_in_wg * 16 + j * 8 + lane_idx / 4], + smem_storage->m_wg[1][warp_idx_in_wg * 16 + j * 8 + lane_idx / 4]); + float o_scale = math::ptx_exp2(m_prev[j] * sm_scale - m[j] * sm_scale); + d[j] *= o_scale; +#pragma unroll + for (uint32_t mma_d = 0; mma_d < KTraits::NUM_MMA_D_CKV / 2; ++mma_d) { + o_frag[mma_d][j * 2 + 0] *= o_scale; + o_frag[mma_d][j * 2 + 1] *= o_scale; + o_frag[mma_d][j * 2 + 4] *= o_scale; + o_frag[mma_d][j * 2 + 5] *= o_scale; + } +#pragma unroll + for (uint32_t mma_kv = 0; mma_kv < KTraits::NUM_MMA_KV / 2; ++mma_kv) { + s_frag[mma_kv][j * 2 + 0] = + math::ptx_exp2(s_frag[mma_kv][j * 2 + 0] * sm_scale - m[j] * sm_scale); + s_frag[mma_kv][j * 2 + 1] = + math::ptx_exp2(s_frag[mma_kv][j * 2 + 1] * sm_scale - m[j] * sm_scale); + s_frag[mma_kv][j * 2 + 4] = + math::ptx_exp2(s_frag[mma_kv][j * 2 + 4] * sm_scale - m[j] * sm_scale); + s_frag[mma_kv][j * 2 + 5] = + math::ptx_exp2(s_frag[mma_kv][j * 2 + 5] * sm_scale - m[j] * sm_scale); + } + } + } else { +#pragma unroll + for (uint32_t j = 0; j < 2; ++j) { + m_prev[j] = m[j]; +#pragma unroll + for (uint32_t mma_kv = 0; mma_kv < KTraits::NUM_MMA_KV; ++mma_kv) { + float m_local = max(max(s_frag[mma_kv][j * 2 + 0], s_frag[mma_kv][j * 2 + 1]), + max(s_frag[mma_kv][j * 2 + 4], s_frag[mma_kv][j * 2 + 5])); + m[j] = max(m[j], m_local); + } + m[j] = max(m[j], math::shfl_xor_sync(m[j], 0x2)); + m[j] = max(m[j], math::shfl_xor_sync(m[j], 0x1)); + } + +#pragma unroll + for (uint32_t j = 0; j < 2; ++j) { + float o_scale = math::ptx_exp2(m_prev[j] * sm_scale - m[j] * sm_scale); + d[j] *= o_scale; +#pragma unroll + for (uint32_t mma_d = 0; mma_d < KTraits::NUM_MMA_D_CKV / 2; ++mma_d) { + o_frag[mma_d][j * 2 + 0] *= o_scale; + o_frag[mma_d][j * 2 + 1] *= o_scale; + o_frag[mma_d][j * 2 + 4] *= o_scale; + o_frag[mma_d][j * 2 + 5] *= o_scale; + } +#pragma unroll + for (uint32_t mma_kv = 0; mma_kv < KTraits::NUM_MMA_KV; ++mma_kv) { + s_frag[mma_kv][j * 2 + 0] = + math::ptx_exp2(s_frag[mma_kv][j * 2 + 0] * sm_scale - m[j] * sm_scale); + s_frag[mma_kv][j * 2 + 1] = + math::ptx_exp2(s_frag[mma_kv][j * 2 + 1] * sm_scale - m[j] * sm_scale); + s_frag[mma_kv][j * 2 + 4] = + math::ptx_exp2(s_frag[mma_kv][j * 2 + 4] * sm_scale - m[j] * sm_scale); + s_frag[mma_kv][j * 2 + 5] = + math::ptx_exp2(s_frag[mma_kv][j * 2 + 5] * sm_scale - m[j] * sm_scale); + } + } + } +} + +template +__device__ __forceinline__ void compute_mla_qk(typename KTraits::SharedStorage* smem_storage, + const uint32_t stage_idx, + typename KTraits::DTypeQKAccum (*s_frag)[8]) { + constexpr uint32_t UPCAST_STRIDE_Q_NOPE = KTraits::UPCAST_STRIDE_Q_NOPE; + constexpr uint32_t UPCAST_STRIDE_Q_PE = KTraits::UPCAST_STRIDE_Q_PE; + constexpr uint32_t UPCAST_STRIDE_CKV = KTraits::UPCAST_STRIDE_CKV; + constexpr uint32_t UPCAST_STRIDE_KPE = KTraits::UPCAST_STRIDE_KPE; + constexpr uint32_t NUM_MMA_KV = KTraits::NUM_MMA_KV; + smem_t q_smem_nope(smem_storage->q_smem_nope); + smem_t q_smem_pe(smem_storage->q_smem_pe); + smem_t ckv_smem(smem_storage->ckv_smem[stage_idx]); + smem_t kpe_smem(smem_storage->kpe_p_smem[stage_idx]); + const uint32_t lane_idx = threadIdx.x, warpgroup_idx = threadIdx.z, warp_idx_in_wg = threadIdx.y; + compute_qk_(q_smem_pe, kpe_smem, s_frag); + compute_qk_(q_smem_nope, ckv_smem, s_frag); +} + +template +__device__ __forceinline__ void compute_mla_pv(typename KTraits::SharedStorage* smem_storage, + const uint32_t stage_idx, + typename KTraits::DTypeQKAccum (*s_frag)[8], + typename KTraits::DTypeQKAccum* d, + float (*o_frag)[8]) { + const uint32_t lane_idx = threadIdx.x, warpgroup_idx = threadIdx.z, warp_idx_in_wg = threadIdx.y; + constexpr uint32_t NUM_MMA_KV = KTraits::NUM_MMA_KV; + constexpr uint32_t NUM_MMA_D_CKV = KTraits::NUM_MMA_D_CKV; + constexpr uint32_t UPCAST_STRIDE_CKV = KTraits::UPCAST_STRIDE_CKV; + smem_t ckv_smem(smem_storage->ckv_smem[stage_idx]); + uint32_t ckv_smem_offset_r = ckv_smem.template get_permuted_offset( + lane_idx % 16, warpgroup_idx * NUM_MMA_D_CKV + lane_idx / 16); + if constexpr (KTraits::QK_SHARD) { + // shard s_frag computation on KV dimension across warpgroups, need allgather + alignas(16) typename KTraits::DTypeKV p_f16[NUM_MMA_KV / 2][8]; +#pragma unroll + for (uint32_t mma_kv = 0; mma_kv < NUM_MMA_KV / 2; ++mma_kv) { + vec_cast::cast<8>(p_f16[mma_kv], s_frag[mma_kv]); + mma::m16k16_rowsum_f16f16f32(d, p_f16[mma_kv]); + } + + __syncthreads(); + smem_t p_smem(smem_storage->kpe_p_smem[stage_idx]); + constexpr uint32_t UPCAST_STRIDE_P = KTraits::UPCAST_STRIDE_P; +#pragma unroll + for (uint32_t mma_kv = 0; mma_kv < NUM_MMA_KV / 2; ++mma_kv) { +#ifdef FLASHINFER_STMATRIX_M8N8X4_ENABLED + uint32_t p_smem_offset_w = p_smem.template get_permuted_offset( + warp_idx_in_wg * 16 + lane_idx % 16, + warpgroup_idx * NUM_MMA_KV + mma_kv * 2 + lane_idx / 16); + p_smem.stmatrix_m8n8x4(p_smem_offset_w, (uint32_t*)p_f16[mma_kv]); +#else + uint32_t p_smem_offset_w = p_smem.template get_permuted_offset( + warp_idx_in_wg * 16 + lane_idx / 4, warpgroup_idx * NUM_MMA_KV + mma_kv * 2); + ((uint32_t*)(p_smem.base + p_smem_offset_w))[lane_idx % 4] = *(uint32_t*)&p_f16[mma_kv][0]; + ((uint32_t*)(p_smem.base + p_smem_offset_w + 8 * UPCAST_STRIDE_P))[lane_idx % 4] = + *(uint32_t*)&p_f16[mma_kv][2]; + ((uint32_t*)(p_smem.base + (p_smem_offset_w ^ 0x1)))[lane_idx % 4] = + *(uint32_t*)&p_f16[mma_kv][4]; + ((uint32_t*)(p_smem.base + (p_smem_offset_w ^ 0x1) + 8 * UPCAST_STRIDE_P))[lane_idx % 4] = + *(uint32_t*)&p_f16[mma_kv][6]; +#endif + } + uint32_t p_smem_offset_r = p_smem.template get_permuted_offset( + warp_idx_in_wg * 16 + lane_idx % 16, lane_idx / 16); + + // wait for p_smem to be filled + __syncthreads(); + +#pragma unroll + for (uint32_t mma_kv = 0; mma_kv < NUM_MMA_KV; ++mma_kv) { + uint32_t p_frag[4]; + p_smem.ldmatrix_m8n8x4(p_smem_offset_r, p_frag); + p_smem_offset_r = p_smem.template advance_offset_by_column<2>(p_smem_offset_r, mma_kv); + +#pragma unroll + for (uint32_t mma_d = 0; mma_d < NUM_MMA_D_CKV / 2; ++mma_d) { + uint32_t v_frag[4]; + ckv_smem.ldmatrix_m8n8x4_trans(ckv_smem_offset_r, v_frag); + mma::mma_sync_m16n16k16_row_col_f16f16f32(o_frag[mma_d], p_frag, + v_frag); + ckv_smem_offset_r = ckv_smem.template advance_offset_by_column<2>(ckv_smem_offset_r, mma_d); + } + ckv_smem_offset_r = + ckv_smem.template advance_offset_by_row<16, UPCAST_STRIDE_CKV>(ckv_smem_offset_r) - + NUM_MMA_D_CKV; + } + } else { + // no need to store p_smem because all warpgroups are working on the same p + alignas(16) typename KTraits::DTypeKV p_f16[NUM_MMA_KV][8]; +#pragma unroll + for (uint32_t mma_kv = 0; mma_kv < NUM_MMA_KV; ++mma_kv) { + vec_cast::cast<8>(p_f16[mma_kv], s_frag[mma_kv]); + mma::m16k16_rowsum_f16f16f32(d, p_f16[mma_kv]); + } +#pragma unroll + for (uint32_t mma_kv = 0; mma_kv < NUM_MMA_KV; ++mma_kv) { +#pragma unroll + for (uint32_t mma_d = 0; mma_d < NUM_MMA_D_CKV / 2; ++mma_d) { + uint32_t v_frag[4]; + ckv_smem.ldmatrix_m8n8x4_trans(ckv_smem_offset_r, v_frag); + mma::mma_sync_m16n16k16_row_col_f16f16f32( + o_frag[mma_d], (uint32_t*)p_f16[mma_kv], v_frag); + ckv_smem_offset_r = ckv_smem.template advance_offset_by_column<2>(ckv_smem_offset_r, mma_d); + } + ckv_smem_offset_r = + ckv_smem.template advance_offset_by_row<16, UPCAST_STRIDE_CKV>(ckv_smem_offset_r) - + NUM_MMA_D_CKV; + } + } +} + +template +__device__ __forceinline__ void normalize_d_(typename KTraits::SharedStorage* smem_storage, + const uint32_t stage_idx, float (*o_frag)[8], + typename KTraits::DTypeQKAccum* m, float* d) { + const uint32_t warpgroup_idx = threadIdx.z, lane_idx = threadIdx.x, warp_idx_in_wg = threadIdx.y; + if constexpr (KTraits::QK_SHARD) { +#pragma unroll + for (uint32_t j = 0; j < 2; ++j) { + if (lane_idx % 4 == 0) { + smem_storage->d_wg[warpgroup_idx][warp_idx_in_wg * 16 + j * 8 + lane_idx / 4] = d[j]; + } + } + __syncthreads(); +#pragma unroll + for (uint32_t j = 0; j < 2; ++j) { + d[j] = smem_storage->d_wg[0][warp_idx_in_wg * 16 + j * 8 + lane_idx / 4] + + smem_storage->d_wg[1][warp_idx_in_wg * 16 + j * 8 + lane_idx / 4]; + } + } + + float d_rcp[2]; + // compute reciprocal of d +#pragma unroll + for (uint32_t j = 0; j < 2; ++j) { + d_rcp[j] = (m[j] != typename KTraits::DTypeQKAccum(-math::inf)) ? math::ptx_rcp(d[j]) : 0.f; + } + +#pragma unroll + for (uint32_t mma_d = 0; mma_d < KTraits::NUM_MMA_D_CKV / 2; ++mma_d) { +#pragma unroll + for (uint32_t reg_id = 0; reg_id < 8; ++reg_id) { + o_frag[mma_d][reg_id] = o_frag[mma_d][reg_id] * d_rcp[(reg_id % 4) / 2]; + } + } +} + +template +__device__ __forceinline__ void finalize_m_(typename KTraits::AttentionVariant variant, + typename KTraits::DTypeQKAccum* m) { + if constexpr (variant.use_softmax) { +#pragma unroll + for (uint32_t j = 0; j < 2; ++j) { + if (m[j] != typename KTraits::DTypeQKAccum(-math::inf)) { + m[j] *= variant.sm_scale_log2; + } + } + } +} + +template +__device__ void DevicePersistentMergeStates( + typename KTraits::IdType* merge_packed_offset_start, + typename KTraits::IdType* merge_packed_offset_end, + typename KTraits::IdType* merge_partial_packed_offset_start, + typename KTraits::IdType* merge_partial_packed_offset_end, + typename KTraits::IdType* merge_partial_stride, typename KTraits::DTypeO* partial_o, + float* partial_lse, typename KTraits::DTypeO* final_o, float* final_lse, + const uint32_t o_stride_n, const uint32_t o_stride_h, const uint_fastdiv& num_heads) { + constexpr uint32_t VEC_SIZE = 8; // partial o has data type float + constexpr uint32_t NUM_THRS_PER_ROW = KTraits::HEAD_DIM_CKV / VEC_SIZE; + constexpr uint32_t ROWS_PER_ITERATION = (KTraits::NUM_THREADS) / NUM_THRS_PER_ROW; + const uint32_t cta_idx = (gridDim.x * blockIdx.y + blockIdx.x); + const uint32_t thread_id = (threadIdx.z * blockDim.y + threadIdx.y) * blockDim.x + threadIdx.x; + const uint32_t offset_start = merge_packed_offset_start[cta_idx]; + const uint32_t len = merge_packed_offset_end[cta_idx] - offset_start; + const uint32_t partial_offset_start = merge_partial_packed_offset_start[cta_idx]; + const uint32_t partial_offset_end = merge_partial_packed_offset_end[cta_idx]; + const uint32_t stride = merge_partial_stride[cta_idx]; +#pragma unroll 1 + for (uint32_t local_packed_offset = thread_id / NUM_THRS_PER_ROW; local_packed_offset < len; + local_packed_offset += ROWS_PER_ITERATION) { + uint32_t final_packed_offset = offset_start + local_packed_offset; + uint32_t q, r; + num_heads.divmod(final_packed_offset, q, r); + state_t st; +#pragma unroll 8 + for (uint32_t partial_packed_offset = partial_offset_start + local_packed_offset; + partial_packed_offset < partial_offset_end; partial_packed_offset += stride) { + vec_t o_partial; + float lse_partial; + o_partial.cast_load(partial_o + partial_packed_offset * KTraits::HEAD_DIM_CKV + + (thread_id % NUM_THRS_PER_ROW) * VEC_SIZE); + lse_partial = partial_lse[partial_packed_offset]; + st.merge(o_partial, lse_partial, 1); + } + st.normalize(); + st.o.cast_store(final_o + + (q * o_stride_n + r * o_stride_h + (thread_id % NUM_THRS_PER_ROW) * VEC_SIZE)); + if (final_lse) { + final_lse[q * num_heads + r] = st.get_lse(); + } + } +} + +template +__device__ __forceinline__ void write_o(typename KTraits::SharedStorage* smem_storage, + typename KTraits::DTypeO* final_o, float* final_lse, + typename KTraits::DTypeO* partial_o, float* partial_lse, + float (*o_frag)[8], typename KTraits::DTypeQKAccum* m, + float* d, const uint32_t o_stride_n, + const uint32_t o_stride_h, const uint32_t q_len, + const uint32_t packed_offset, + const uint_fastdiv& num_heads) { + using DTypeO = typename KTraits::DTypeO; + constexpr uint32_t NUM_MMA_D_CKV = KTraits::NUM_MMA_D_CKV; + constexpr uint32_t HEAD_DIM_CKV = KTraits::HEAD_DIM_CKV; + constexpr uint32_t UPCAST_STRIDE_FINAL_O = KTraits::UPCAST_STRIDE_FINAL_O; + const uint32_t lane_idx = threadIdx.x, warpgroup_idx = threadIdx.z, warp_idx_in_wg = threadIdx.y; + smem_t o_smem(smem_storage->o_smem); +#pragma unroll + for (uint32_t mma_d = 0; mma_d < NUM_MMA_D_CKV / 2; ++mma_d) { + uint32_t o_frag_f16[8 / 2]; + vec_cast::cast<8>((DTypeO*)o_frag_f16, o_frag[mma_d]); +#ifdef FLASHINFER_STMATRIX_M8N8X4_ENABLED + uint32_t o_smem_offset_w = o_smem.template get_permuted_offset( + warp_idx_in_wg * 16 + lane_idx % 16, + warpgroup_idx * NUM_MMA_D_CKV + mma_d * 2 + lane_idx / 16); + o_smem.template stmatrix_m8n8x4(o_smem_offset_w, o_frag_f16); +#else + uint32_t o_smem_offset_w = o_smem.template get_permuted_offset( + warp_idx_in_wg * 16 + lane_idx / 4, warpgroup_idx * NUM_MMA_D_CKV + mma_d * 2); + ((uint32_t*)(o_smem.base + o_smem_offset_w))[lane_idx % 4] = o_frag_f16[0]; + ((uint32_t*)(o_smem.base + o_smem_offset_w + 8 * UPCAST_STRIDE_FINAL_O))[lane_idx % 4] = + o_frag_f16[1]; + ((uint32_t*)(o_smem.base + (o_smem_offset_w ^ 0x1)))[lane_idx % 4] = o_frag_f16[2]; + ((uint32_t*)(o_smem.base + (o_smem_offset_w ^ 0x1) + 8 * UPCAST_STRIDE_FINAL_O))[lane_idx % 4] = + o_frag_f16[3]; +#endif + } + + if (partial_o != nullptr) { + // write to partial_o +#pragma unroll + for (uint32_t j = 0; j < 2; ++j) { + uint32_t q_idx = (packed_offset + warp_idx_in_wg * 16 + 8 * j + lane_idx / 4) / num_heads; + if (lane_idx % 4 == 0 && q_idx < q_len) { + partial_lse[(blockIdx.x * 4 + warp_idx_in_wg) * 16 + 8 * j + lane_idx / 4] = + math::ptx_log2(d[j]) + float(m[j]); + } + } + + // step 1. smem to gmem + uint32_t o_smem_offset_w = o_smem.template get_permuted_offset( + warp_idx_in_wg * 16 + lane_idx / 8, warpgroup_idx * NUM_MMA_D_CKV + lane_idx % 8); +#pragma unroll + for (uint32_t j = 0; j < 4; ++j) { + uint32_t q_idx = (packed_offset + warp_idx_in_wg * 16 + 4 * j + lane_idx / 8) / num_heads; + DTypeO* o_partial_ptr = + partial_o + + ((blockIdx.x * 4 + warp_idx_in_wg) * 16 + 4 * j + lane_idx / 8) * HEAD_DIM_CKV + + warpgroup_idx * (HEAD_DIM_CKV / 2) + (lane_idx % 8) * upcast_size(); +#pragma unroll + for (uint32_t mma_d = 0; mma_d < NUM_MMA_D_CKV / 8; ++mma_d) { + if (q_idx < q_len) { + o_smem.template store_128b(o_smem_offset_w, o_partial_ptr); + } + o_partial_ptr += 8 * upcast_size(); + o_smem_offset_w = o_smem.template advance_offset_by_column<8>(o_smem_offset_w, mma_d); + } + o_smem_offset_w = + o_smem.template advance_offset_by_row<4, UPCAST_STRIDE_FINAL_O>(o_smem_offset_w) - + NUM_MMA_D_CKV; + } + } else { + // write to final_o + + if (final_lse) { +#pragma unroll + for (uint32_t j = 0; j < 2; ++j) { + uint32_t q, r; + num_heads.divmod(packed_offset + warp_idx_in_wg * 16 + 8 * j + lane_idx / 4, q, r); + if (lane_idx % 4 == 0 && q < q_len) { + final_lse[q * num_heads + r] = math::ptx_log2(d[j]) + float(m[j]); + } + } + } + + // step 1. smem to gmem + uint32_t o_smem_offset_w = o_smem.template get_permuted_offset( + warp_idx_in_wg * 16 + lane_idx / 8, warpgroup_idx * NUM_MMA_D_CKV + lane_idx % 8); +#pragma unroll + for (uint32_t j = 0; j < 4; ++j) { + uint32_t q, r; + num_heads.divmod(packed_offset + warp_idx_in_wg * 16 + 4 * j + lane_idx / 8, q, r); + DTypeO* o_final_ptr = final_o + q * o_stride_n + r * o_stride_h + + warpgroup_idx * (HEAD_DIM_CKV / 2) + + (lane_idx % 8) * upcast_size(); +#pragma unroll + for (uint32_t mma_d = 0; mma_d < NUM_MMA_D_CKV / 8; ++mma_d) { + if (q < q_len) { + o_smem.template store_128b(o_smem_offset_w, o_final_ptr); + } + o_final_ptr += 8 * upcast_size(); + o_smem_offset_w = o_smem.template advance_offset_by_column<8>(o_smem_offset_w, mma_d); + } + o_smem_offset_w = + o_smem.template advance_offset_by_row<4, UPCAST_STRIDE_FINAL_O>(o_smem_offset_w) - + NUM_MMA_D_CKV; + } + } +} + +template +__global__ __launch_bounds__(KTraits::NUM_THREADS) void BatchMLAPagedAttentionKernel( + const __grid_constant__ Params params) { + using DTypeQ = typename Params::DTypeQ; + using DTypeKV = typename Params::DTypeKV; + using DTypeO = typename Params::DTypeO; + using IdType = typename Params::IdType; + + extern __shared__ __align__(alignof(typename KTraits::SharedStorage)) uint8_t smem[]; + auto& smem_storage = reinterpret_cast(smem); + + typename KTraits::AttentionVariant variant(params, blockIdx.y, smem); + + [[maybe_unused]] constexpr SwizzleMode SWIZZLE_MODE_Q_NOPE = KTraits::SWIZZLE_MODE_Q_NOPE; + [[maybe_unused]] constexpr SwizzleMode SWIZZLE_MODE_Q_PE = KTraits::SWIZZLE_MODE_Q_PE; + [[maybe_unused]] constexpr SwizzleMode SWIZZLE_MODE_CKV = KTraits::SWIZZLE_MODE_CKV; + [[maybe_unused]] constexpr SwizzleMode SWIZZLE_MODE_KPE = KTraits::SWIZZLE_MODE_KPE; + [[maybe_unused]] constexpr uint32_t NUM_MMA_KV = KTraits::NUM_MMA_KV; + [[maybe_unused]] constexpr uint32_t NUM_MMA_D_CKV = KTraits::NUM_MMA_D_CKV; + [[maybe_unused]] constexpr uint32_t CTA_TILE_Q = KTraits::CTA_TILE_Q; + [[maybe_unused]] constexpr uint32_t CTA_TILE_KV = KTraits::CTA_TILE_KV; + [[maybe_unused]] constexpr int32_t NUM_STAGES = KTraits::NUM_STAGES; + [[maybe_unused]] constexpr bool CAUSAL = KTraits::CAUSAL; + + DTypeQ* q_nope = params.q_nope; + DTypeQ* q_pe = params.q_pe; + DTypeKV* ckv = params.ckv; + DTypeKV* kpe = params.kpe; + IdType* kv_indices = params.kv_indices; + DTypeO* partial_o = params.partial_o; + float* partial_lse = params.partial_lse; + DTypeO* final_o = params.final_o; + float* final_lse = params.final_lse; + IdType* work_indptr = params.work_indptr; + + float s_frag[KTraits::QK_SHARD ? NUM_MMA_KV / 2 : NUM_MMA_KV][8]; + alignas(16) float o_frag[NUM_MMA_D_CKV / 2][8]; + float m[2]; + float d[2]; + + const uint_fastdiv& num_heads = params.num_heads; + const uint_fastdiv& block_size = params.block_size; + const uint32_t q_nope_stride_n = params.q_nope_stride_n; + const uint32_t q_nope_stride_h = params.q_nope_stride_h; + const uint32_t q_pe_stride_n = params.q_pe_stride_n; + const uint32_t q_pe_stride_h = params.q_pe_stride_h; + const uint32_t ckv_stride_page = params.ckv_stride_page; + const uint32_t ckv_stride_n = params.ckv_stride_n; + const uint32_t kpe_stride_page = params.kpe_stride_page; + const uint32_t kpe_stride_n = params.kpe_stride_n; + const uint32_t o_stride_n = params.o_stride_n; + const uint32_t o_stride_h = params.o_stride_h; + const uint32_t cluster_tile_q = gridDim.x * KTraits::CTA_TILE_Q; + +#pragma unroll 1 + for (IdType work_idx = work_indptr[blockIdx.y]; work_idx < work_indptr[blockIdx.y + 1]; + ++work_idx) { + const uint32_t q_indptr = params.q_indptr[work_idx]; + const uint32_t kv_indptr = params.kv_indptr[work_idx]; + const int32_t partial_indptr = params.partial_indptr[work_idx]; + const uint32_t q_len = params.q_len[work_idx]; + const uint32_t kv_len = params.kv_len[work_idx]; + const uint32_t packed_qo_start = params.q_start[work_idx]; + const uint32_t kv_start = params.kv_start[work_idx]; + const uint32_t kv_end = params.kv_end[work_idx]; + + const uint32_t qo_packed_idx_base = packed_qo_start + blockIdx.x * KTraits::CTA_TILE_Q; + const uint32_t qo_upperbound = + min(q_len, ceil_div(qo_packed_idx_base + KTraits::CTA_TILE_Q, num_heads)); + + init_states_(o_frag, m, d); + + __syncthreads(); + load_q(&smem_storage, q_nope + q_indptr * q_nope_stride_n, + q_pe + q_indptr * q_pe_stride_n, q_nope_stride_n, q_nope_stride_h, + q_pe_stride_n, q_pe_stride_h, qo_upperbound, qo_packed_idx_base, + params.num_heads); + + int kv_tile_idx = + ceil_div( + (CAUSAL ? min(kv_end, kv_len - q_len + (packed_qo_start + cluster_tile_q) / num_heads) + : kv_end), + CTA_TILE_KV) - + 1 - (kv_start / CTA_TILE_KV); + + int mask_tile_idx = + (CAUSAL ? min(kv_end, kv_len - q_len + packed_qo_start / num_heads) : kv_end) / + CTA_TILE_KV - + (kv_start / CTA_TILE_KV); + + uint32_t block_iter_base = kv_indptr * block_size + kv_start; + // last kv tile + __syncthreads(); + uint32_t packed_kv_bound = kv_indptr * block_size + kv_len; + load_kv(&smem_storage, ckv, kpe, kv_indices, ckv_stride_n, ckv_stride_page, + kpe_stride_n, kpe_stride_page, packed_kv_bound, + block_iter_base + kv_tile_idx * CTA_TILE_KV, block_size, + kv_tile_idx % NUM_STAGES); + cp_async::commit_group(); +#pragma unroll + for (int stage_idx = 1; stage_idx < NUM_STAGES; ++stage_idx) { + if (kv_tile_idx - stage_idx >= 0) { + load_kv(&smem_storage, ckv, kpe, kv_indices, ckv_stride_n, ckv_stride_page, + kpe_stride_n, kpe_stride_page, packed_kv_bound, + block_iter_base + (kv_tile_idx - stage_idx) * CTA_TILE_KV, block_size, + (kv_tile_idx - stage_idx) % NUM_STAGES); + cp_async::commit_group(); + } + } + + // loop with mask +#pragma unroll 1 + for (; kv_tile_idx >= mask_tile_idx && kv_tile_idx > 0; --kv_tile_idx) { + cp_async::wait_group(); + __syncthreads(); + + // compute mla qk + compute_mla_qk(&smem_storage, kv_tile_idx % NUM_STAGES, s_frag); + + // logits mask + logits_mask_(qo_packed_idx_base, kv_start + kv_tile_idx * CTA_TILE_KV, q_len, kv_len, + kv_end, num_heads, s_frag); + + // compute m,d states in online softmax + update_mdo_states_(&smem_storage, kv_tile_idx % NUM_STAGES, variant, s_frag, o_frag, + m, d); + + // compute sfm * v + compute_mla_pv(&smem_storage, kv_tile_idx % NUM_STAGES, s_frag, d, o_frag); + + if (kv_tile_idx - NUM_STAGES >= 0) { + __syncthreads(); + load_kv(&smem_storage, ckv, kpe, kv_indices, ckv_stride_n, ckv_stride_page, + kpe_stride_n, kpe_stride_page, packed_kv_bound, + block_iter_base + (kv_tile_idx - NUM_STAGES) * CTA_TILE_KV, block_size, + (kv_tile_idx - NUM_STAGES) % NUM_STAGES); + cp_async::commit_group(); + } + } + + // loop without mask +#pragma unroll 1 + for (; kv_tile_idx + 1 > NUM_STAGES; --kv_tile_idx) { + cp_async::wait_group(); + __syncthreads(); + + // compute mla qk + compute_mla_qk(&smem_storage, kv_tile_idx % NUM_STAGES, s_frag); + + // compute m,d states in online softmax + update_mdo_states_(&smem_storage, kv_tile_idx % NUM_STAGES, variant, s_frag, o_frag, + m, d); + // compute sfm * v + compute_mla_pv(&smem_storage, kv_tile_idx % NUM_STAGES, s_frag, d, o_frag); + + __syncthreads(); + load_kv(&smem_storage, ckv, kpe, kv_indices, ckv_stride_n, ckv_stride_page, + kpe_stride_n, kpe_stride_page, packed_kv_bound, + block_iter_base + (kv_tile_idx - NUM_STAGES) * CTA_TILE_KV, block_size, + (kv_tile_idx - NUM_STAGES) % NUM_STAGES); + cp_async::commit_group(); + } + cp_async::wait_group<0>(); + __syncthreads(); + + // last tiles +#pragma unroll + for (; kv_tile_idx >= 0; --kv_tile_idx) { + // compute mla qk + compute_mla_qk(&smem_storage, kv_tile_idx % NUM_STAGES, s_frag); + + logits_mask_(qo_packed_idx_base, kv_start + kv_tile_idx * CTA_TILE_KV, q_len, kv_len, + kv_end, num_heads, s_frag); + + // compute m,d states in online softmax + update_mdo_states_(&smem_storage, kv_tile_idx % NUM_STAGES, variant, s_frag, o_frag, + m, d); + + // compute sfm * v + compute_mla_pv(&smem_storage, kv_tile_idx % NUM_STAGES, s_frag, d, o_frag); + } + + __syncthreads(); + + // normalize and write back + normalize_d_(&smem_storage, kv_tile_idx % NUM_STAGES, o_frag, m, d); + + finalize_m_(variant, m); + + write_o( + &smem_storage, final_o + q_indptr * o_stride_n, + final_lse ? final_lse + q_indptr * num_heads : nullptr, + (partial_indptr == -1) ? nullptr : partial_o + partial_indptr * KTraits::HEAD_DIM_CKV, + (partial_indptr == -1) ? nullptr : partial_lse + partial_indptr, o_frag, m, d, o_stride_n, + o_stride_h, qo_upperbound, qo_packed_idx_base, num_heads); + } + + auto grid = cg::this_grid(); + grid.sync(); + + // the second stage, merge partial outputs + DevicePersistentMergeStates( + params.merge_packed_offset_start, params.merge_packed_offset_end, + params.merge_partial_packed_offset_start, params.merge_partial_packed_offset_end, + params.merge_partial_stride, partial_o, partial_lse, final_o, final_lse, o_stride_n, + o_stride_h, num_heads); +} + +#define DISPATCH_SMEM_CONFIG(smem_limit_per_sm, NUM_STAGES, CTA_TILE_KV, QK_SHARD, ...) \ + if (smem_limit_per_sm >= 221696) { \ + constexpr uint32_t NUM_STAGES = 2; \ + constexpr uint32_t CTA_TILE_KV = 64; \ + constexpr bool QK_SHARD = true; \ + __VA_ARGS__; \ + } else if (smem_limit_per_sm >= 147968) { \ + constexpr uint32_t NUM_STAGES = 2; \ + constexpr uint32_t CTA_TILE_KV = 32; \ + constexpr bool QK_SHARD = true; \ + __VA_ARGS__; \ + } else if (smem_limit_per_sm >= 92672) { \ + constexpr uint32_t NUM_STAGES = 1; \ + constexpr uint32_t CTA_TILE_KV = 16; \ + constexpr bool QK_SHARD = false; \ + __VA_ARGS__; \ + } else { \ + std::ostringstream err; \ + err << "Unsupported shared memory size: " << smem_limit_per_sm; \ + FLASHINFER_ERROR(err.str()); \ + return cudaErrorNotSupported; \ + } + +template +cudaError_t BatchMLAPagedAttention(Params params, uint32_t num_blks_x, uint32_t num_blks_y, + cudaStream_t stream) { + using DTypeQ = typename Params::DTypeQ; + using DTypeKV = typename Params::DTypeKV; + using DTypeO = typename Params::DTypeO; + using IdType = typename Params::IdType; + if (MASK_MODE == MaskMode::kCustom) { + return cudaErrorNotSupported; + } + constexpr bool CAUSAL = MASK_MODE == MaskMode::kCausal; + + dim3 nblks(num_blks_x, num_blks_y); + dim3 nthrs(32, 4, 2); + + // get GPU shared memory size + int device; + int smem_limit_per_sm; + cudaGetDevice(&device); + cudaDeviceGetAttribute(&smem_limit_per_sm, cudaDevAttrMaxSharedMemoryPerMultiprocessor, device); + + DISPATCH_SMEM_CONFIG(smem_limit_per_sm, NUM_STAGES, CTA_TILE_KV, QK_SHARD, { + using KTraits = KernelTraits; + size_t smem_size = sizeof(typename KTraits::SharedStorage); + auto kernel = BatchMLAPagedAttentionKernel; + void* args[] = {(void*)¶ms}; + + FLASHINFER_CUDA_CALL( + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + FLASHINFER_CUDA_CALL( + cudaLaunchCooperativeKernel((void*)kernel, nblks, nthrs, args, smem_size, stream)); + }); + + return cudaSuccess; +} + +} // namespace mla + +} // namespace flashinfer + +#endif // FLASHINFER_MLA_FA2_CUH_ diff --git a/include/flashinfer/attention/mla_hopper.cuh b/include/flashinfer/attention/mla_hopper.cuh new file mode 100644 index 0000000000000000000000000000000000000000..f76c110dd899cd46d936adb41aad093cbb27fbd5 --- /dev/null +++ b/include/flashinfer/attention/mla_hopper.cuh @@ -0,0 +1,1007 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef FLASHINFER_MLA_HOPPER_CUH_ +#define FLASHINFER_MLA_HOPPER_CUH_ +#include + +#include +#include + +#include "hopper.cuh" +#include "mla.cuh" +#include "mla_params.cuh" +#include "prefill.cuh" +#include "variant_helper.cuh" + +namespace flashinfer { + +namespace mla { + +namespace hopper { + +enum class ProfileEventType { + kIssueLoadQ = 0U, + kIssueLoadKV = 1U, + kWriteO = 2U, + kSoftmaxUpdate = 3U, + kGemmQK = 4U, + kGemmPV = 5U, + kRescaleO = 6U, + kWritePSmem = 7U, + kSplitK = 8U, +}; + +enum class NamedBarriers { kOScaleReady = 1U, kBarrierO = 2U, kMDReady = 3U }; + +__device__ __forceinline__ void barrier_arrive(int num_threads, NamedBarriers barrier) { + cutlass::arch::NamedBarrier::arrive(num_threads, static_cast(barrier)); +} + +__device__ __forceinline__ void barrier_sync(int num_threads, NamedBarriers barrier) { + cutlass::arch::NamedBarrier::sync(num_threads, static_cast(barrier)); +} + +template +struct HopperSharedStorageQKVO { + struct { + struct { + struct { + alignas(16) DTypeQ nope[CTA_TILE_Q * HEAD_DIM_CKV]; + alignas(16) DTypeQ pe[CTA_TILE_Q * HEAD_DIM_KPE]; + } q_smem; + union { + struct { + alignas(16) DTypeKV ckv[CTA_TILE_KV * HEAD_DIM_CKV]; + union { + alignas(16) DTypeKV kpe[CTA_TILE_KV * HEAD_DIM_KPE]; + alignas(16) DTypeKV p[CTA_TILE_Q * CTA_TILE_KV]; + }; + }; + alignas(16) DTypeO o[CTA_TILE_Q * HEAD_DIM_CKV]; + } kv_o_smem[NUM_STAGES]; + alignas(16) float o_scale[CTA_TILE_Q]; + alignas(16) float m[CTA_TILE_Q]; + alignas(16) float d[CTA_TILE_Q]; + }; + + typename MainloopPipeline::SharedStorage pipeline_q, pipeline_kv; + }; +}; + +template +struct HopperKernelTraits + : KernelTraits { + static constexpr uint32_t NUM_THREADS = 256; + static constexpr uint32_t NUM_COPY_THREADS = 128; + static constexpr uint32_t NUM_QK_THREADS = 128; + static constexpr uint32_t NUM_REGS_S_FRAG = CTA_TILE_KV_ / 2; + static constexpr uint32_t NUM_REGS_O_FRAG = HEAD_DIM_CKV_ / 4; + static constexpr uint32_t NUM_REGS_P_FRAG = CTA_TILE_KV_ / 4; + using MainloopPipeline = cutlass::PipelineAsync; + using SharedStorage = + HopperSharedStorageQKVO; +}; + +template +__device__ __forceinline__ void init_states_(float* o_frag, float* m, float* d, float* o_scale) { +#pragma unroll + for (uint32_t reg_id = 0; reg_id < KTraits::NUM_REGS_O_FRAG; ++reg_id) { + o_frag[reg_id] = 0.f; + } + +#pragma unroll + for (uint32_t j = 0; j < 2; ++j) { + m[j] = -math::inf; + d[j] = 1.f; + o_scale[j] = 1.f; + } +} + +template +__device__ __forceinline__ void load_q( + typename KTraits::SharedStorage* smem_storage, typename KTraits::DTypeQ* q_nope, + typename KTraits::DTypeQ* q_pe, const uint32_t q_nope_stride_n, const uint32_t q_nope_stride_h, + const uint32_t q_pe_stride_n, const uint32_t q_pe_stride_h, const uint32_t q_len, + const uint32_t packed_offset, const uint_fastdiv& num_heads) { + using DTypeQ = typename KTraits::DTypeQ; + constexpr uint32_t UPCAST_STRIDE_Q_NOPE = KTraits::UPCAST_STRIDE_Q_NOPE; + constexpr uint32_t UPCAST_STRIDE_Q_PE = KTraits::UPCAST_STRIDE_Q_PE; + constexpr uint32_t NUM_MMA_D_CKV = KTraits::NUM_MMA_D_CKV; + constexpr uint32_t NUM_MMA_D_KPE = KTraits::NUM_MMA_D_KPE; + const uint32_t lane_idx = cutlass::canonical_lane_idx(); + const uint32_t warp_group_idx = cutlass::canonical_warp_group_idx(); + const uint32_t warp_idx_in_wg = cutlass::canonical_warp_idx() % 4; + + smem_t q_smem_nope(smem_storage->q_smem.nope); + smem_t q_smem_pe(smem_storage->q_smem.pe); + +#pragma unroll + for (uint32_t mma_q = 0; mma_q < 2; ++mma_q) { +#pragma unroll + for (uint32_t j = 0; j < 2; ++j) { + uint32_t q, r; + num_heads.divmod(packed_offset + lane_idx / 8 + (j + mma_q * 2) * 16 + warp_idx_in_wg * 4, q, + r); + DTypeQ* q_nope_ptr = q_nope + q * q_nope_stride_n + r * q_nope_stride_h + + (lane_idx % 8) * upcast_size(); + DTypeQ* q_pe_ptr = + q_pe + q * q_pe_stride_n + r * q_pe_stride_h + (lane_idx % 8) * upcast_size(); + uint32_t q_smem_nope_offset_w = + get_swizzle_offset( + 32 * mma_q + j * 16 + warp_idx_in_wg * 4 + lane_idx / 8, 8 * 0 + lane_idx % 8); + uint32_t q_smem_pe_offset_w = + get_swizzle_offset( + 32 * mma_q + j * 16 + warp_idx_in_wg * 4 + lane_idx / 8, 8 * 0 + lane_idx % 8); + +#pragma unroll + for (uint32_t mma_d = 0; mma_d < KTraits::NUM_MMA_D_CKV / 4; ++mma_d) { + q_smem_nope.load_128b_async(q_smem_nope_offset_w, q_nope_ptr, + q < q_len); + q_smem_nope_offset_w += 64; + q_nope_ptr += 8 * upcast_size(); + } +#pragma unroll + for (uint32_t mma_d = 0; mma_d < KTraits::NUM_MMA_D_KPE / 4; ++mma_d) { + q_smem_pe.load_128b_async(q_smem_pe_offset_w, q_pe_ptr, + q < q_len); + q_smem_pe_offset_w += 64; + q_pe_ptr += 8 * upcast_size(); + } + } + } +} + +template +__device__ __forceinline__ void prefetch_offset( + const uint32_t packed_block_iter_base, const uint32_t packed_kv_bound, + const uint32_t ckv_stride_page, const uint32_t ckv_stride_n, const uint32_t kpe_stride_page, + const uint32_t kpe_stride_n, const uint_fastdiv& block_size, typename KTraits::IdType* indices, + int64_t (*ckv_offset)[2], int64_t (*kpe_offset)[2]) { + using DTypeKV = typename KTraits::DTypeKV; + const uint32_t lane_idx = cutlass::canonical_lane_idx(); + const uint32_t warp_idx_in_wg = cutlass::canonical_warp_idx() % 4; +#pragma unroll + for (uint32_t mma_kv = 0; mma_kv < KTraits::NUM_MMA_KV / 2; ++mma_kv) { +#pragma unroll + for (uint32_t j = 0; j < 2; ++j) { + uint32_t q, r; + uint32_t packed_block_iter = + packed_block_iter_base + lane_idx / 8 + (j + mma_kv * 2) * 16 + warp_idx_in_wg * 4; + block_size.divmod(packed_block_iter, q, r); + ckv_offset[mma_kv][j] = + (packed_block_iter < packed_kv_bound ? indices[q] : 0) * ckv_stride_page + + r * ckv_stride_n + (lane_idx % 8) * upcast_size(); + kpe_offset[mma_kv][j] = + (packed_block_iter < packed_kv_bound ? indices[q] : 0) * kpe_stride_page + + r * kpe_stride_n + (lane_idx % 8) * upcast_size(); + } + } +} + +template +__device__ __forceinline__ void load_kv(typename KTraits::SharedStorage* smem_storage, + typename KTraits::DTypeKV* ckv, + typename KTraits::DTypeKV* kpe, + const uint32_t packed_kv_bound, + const uint32_t packed_block_iter_base, + const uint32_t stage_idx, int64_t (*ckv_offset)[2], + int64_t (*kpe_offset)[2]) { + using DTypeKV = typename KTraits::DTypeKV; + constexpr uint32_t UPCAST_STRIDE_CKV = KTraits::UPCAST_STRIDE_CKV; + constexpr uint32_t UPCAST_STRIDE_KPE = KTraits::UPCAST_STRIDE_KPE; + constexpr uint32_t NUM_MMA_D_CKV = KTraits::NUM_MMA_D_CKV; + constexpr uint32_t NUM_MMA_D_KPE = KTraits::NUM_MMA_D_KPE; + const uint32_t lane_idx = cutlass::canonical_lane_idx(); + const uint32_t warp_idx_in_wg = cutlass::canonical_warp_idx() % 4; + + smem_t ckv_smem(smem_storage->kv_o_smem[stage_idx].ckv); + smem_t kpe_smem(smem_storage->kv_o_smem[stage_idx].kpe); + +#pragma unroll + for (uint32_t mma_kv = 0; mma_kv < KTraits::NUM_MMA_KV / 2; ++mma_kv) { +#pragma unroll + for (uint32_t j = 0; j < 2; ++j) { + uint32_t packed_block_iter = + packed_block_iter_base + lane_idx / 8 + (j + mma_kv * 2) * 16 + warp_idx_in_wg * 4; + + DTypeKV* ckv_ptr = ckv + ckv_offset[mma_kv][j]; + DTypeKV* kpe_ptr = kpe + kpe_offset[mma_kv][j]; + uint32_t ckv_smem_offset_w = get_swizzle_offset( + 32 * mma_kv + j * 16 + warp_idx_in_wg * 4 + lane_idx / 8, 8 * 0 + lane_idx % 8); + uint32_t kpe_smem_offset_w = get_swizzle_offset( + 32 * mma_kv + j * 16 + warp_idx_in_wg * 4 + lane_idx / 8, 8 * 0 + lane_idx % 8); + +#pragma unroll + for (uint32_t mma_d = 0; mma_d < KTraits::NUM_MMA_D_CKV / 4; ++mma_d) { + if constexpr (predicate) { + ckv_smem.load_128b_async( + ckv_smem_offset_w, ckv_ptr, packed_block_iter < packed_kv_bound); + } else { + ckv_smem.load_128b_async(ckv_smem_offset_w, ckv_ptr); + } + ckv_smem_offset_w += 64; + ckv_ptr += 8 * upcast_size(); + } + +#pragma unroll + for (uint32_t mma_d = 0; mma_d < KTraits::NUM_MMA_D_KPE / 4; ++mma_d) { + if constexpr (predicate) { + kpe_smem.load_128b_async( + kpe_smem_offset_w, kpe_ptr, packed_block_iter < packed_kv_bound); + } else { + kpe_smem.load_128b_async(kpe_smem_offset_w, kpe_ptr); + } + kpe_smem_offset_w += 64; + kpe_ptr += 8 * upcast_size(); + } + } + } +} + +template +__device__ __forceinline__ void compute_mla_qk(typename KTraits::SharedStorage* smem_storage, + const uint32_t stage_idx, float* s_frag) { + auto desc_q_pe = + make_smem_desc( + smem_storage->q_smem.pe); + auto desc_k_pe = + make_smem_desc( + smem_storage->kv_o_smem[stage_idx].kpe); + using wgmma = WGMMA_ASYNC_SS; + + warpgroup_fence_frag(s_frag); + warpgroup_arrive(); +#pragma unroll + for (uint32_t mma_d_pe = 0; mma_d_pe < KTraits::NUM_MMA_D_KPE; ++mma_d_pe) { + if (mma_d_pe == 0) { + wgmma::op(desc_q_pe, desc_k_pe, s_frag); + } else { + wgmma::op(desc_q_pe, desc_k_pe, s_frag); + } + if ((mma_d_pe + 1) % 4 == 0) { + desc_q_pe += 64 - 6; + desc_k_pe += 64 - 6; + } else { + desc_q_pe += 2; + desc_k_pe += 2; + } + } + + auto desc_q_nope = + make_smem_desc( + smem_storage->q_smem.nope); + auto desc_ckv = + make_smem_desc( + smem_storage->kv_o_smem[stage_idx].ckv); + +#pragma unroll + for (uint32_t mma_d_ckv = 0; mma_d_ckv < KTraits::NUM_MMA_D_CKV; ++mma_d_ckv) { + wgmma::op(desc_q_nope, desc_ckv, s_frag); + if ((mma_d_ckv + 1) % 4 == 0) { + desc_q_nope += 64 - 6; + desc_ckv += 64 - 6; + } else { + desc_q_nope += 2; + desc_ckv += 2; + } + } + + warpgroup_commit_batch(); + warpgroup_fence_frag(s_frag); +} + +template +__device__ __forceinline__ void compute_mla_pv(typename KTraits::SharedStorage* smem_storage, + const uint32_t stage_idx, float* o_frag) { + const uint32_t lane_idx = cutlass::canonical_lane_idx(); + const uint32_t warp_idx_in_wg = cutlass::canonical_warp_idx() % 4; + const uint32_t warp_group_idx = cutlass::canonical_warp_group_idx(); + + auto desc_p = make_smem_desc( + smem_storage->kv_o_smem[stage_idx].p); + auto desc_ckv = + make_smem_desc( + smem_storage->kv_o_smem[stage_idx].ckv + + warp_group_idx * 8 * (KTraits::HEAD_DIM_CKV / 2)); + warpgroup_fence_frag(o_frag); + warpgroup_arrive(); + using wgmma = WGMMA_ASYNC_SS; + +#pragma unroll + for (uint32_t mma_kv = 0; mma_kv < KTraits::NUM_MMA_KV; ++mma_kv) { + wgmma::op(desc_p, desc_ckv, o_frag); + desc_p += 2; + desc_ckv += 1024; + } + warpgroup_commit_batch(); + warpgroup_fence_frag(o_frag); +} + +template +__device__ __forceinline__ void logits_mask_(const uint32_t qo_packed_idx_base, + const uint32_t kv_idx_base, const uint32_t qo_len, + const uint32_t kv_len, const uint32_t kv_end, + const uint_fastdiv num_heads, float* s_frag) { + const uint32_t lane_idx = cutlass::canonical_lane_idx(); + const uint32_t warp_idx_in_wg = cutlass::canonical_warp_idx() % 4; + constexpr uint32_t NUM_MMA_KV = KTraits::NUM_MMA_KV; + uint32_t q[2]; +#pragma unroll + for (uint32_t j = 0; j < 2; ++j) { + q[j] = (qo_packed_idx_base + warp_idx_in_wg * 16 + lane_idx / 4 + 8 * j) / num_heads; + } + +#pragma unroll + for (uint32_t reg_id = 0; reg_id < KTraits::NUM_REGS_S_FRAG; ++reg_id) { + const uint32_t q_idx = q[(reg_id % 4) / 2], + kv_idx = kv_idx_base + 2 * (lane_idx % 4) + 8 * (reg_id / 4) + reg_id % 2; + const bool mask = (!(KTraits::CAUSAL ? (kv_idx + qo_len > kv_len + q_idx || (kv_idx >= kv_end)) + : kv_idx >= kv_end)); + s_frag[reg_id] = (mask) ? s_frag[reg_id] : (KTraits::MaskFillValue); + } +} + +template +__device__ __forceinline__ void rescale_o_(float* o_scale, float* o_frag) { + const uint32_t lane_idx = cutlass::canonical_lane_idx(); + const uint32_t warp_idx_in_wg = cutlass::canonical_warp_idx() % 4; +#pragma unroll + for (uint32_t reg_id = 0; reg_id < KTraits::NUM_REGS_O_FRAG; ++reg_id) { + o_frag[reg_id] *= o_scale[(reg_id % 4) / 2]; + } +} + +template +__device__ __forceinline__ void update_md_(typename KTraits::SharedStorage* smem_storage, + typename KTraits::AttentionVariant variant, + float* s_frag, float* m, float* d, float* o_scale) { + using AttentionVariant = typename KTraits::AttentionVariant; + const float sm_scale = variant.sm_scale_log2; + const uint32_t lane_idx = cutlass::canonical_lane_idx(); + const uint32_t warp_idx_in_wg = cutlass::canonical_warp_idx() % 4; + float m_prev[2]; +#pragma unroll + for (uint32_t j = 0; j < 2; ++j) { + m_prev[j] = m[j]; +#pragma unroll + for (uint32_t k = 0; k < KTraits::NUM_REGS_S_FRAG / 4; ++k) { + float m_local = max(s_frag[k * 4 + j * 2 + 0], s_frag[k * 4 + j * 2 + 1]); + m[j] = max(m[j], m_local); + } + m[j] = max(m[j], math::shfl_xor_sync(m[j], 0x2)); + m[j] = max(m[j], math::shfl_xor_sync(m[j], 0x1)); + } + +#pragma unroll + for (uint32_t j = 0; j < 2; ++j) { + o_scale[j] = math::ptx_exp2(m_prev[j] * sm_scale - m[j] * sm_scale); + float d_local = 0.f; +#pragma unroll + for (uint32_t k = 0; k < KTraits::NUM_REGS_S_FRAG / 4; ++k) { + s_frag[k * 4 + j * 2 + 0] = + math::ptx_exp2(s_frag[k * 4 + j * 2 + 0] * sm_scale - m[j] * sm_scale); + s_frag[k * 4 + j * 2 + 1] = + math::ptx_exp2(s_frag[k * 4 + j * 2 + 1] * sm_scale - m[j] * sm_scale); + + d_local += s_frag[k * 4 + j * 2 + 0] + s_frag[k * 4 + j * 2 + 1]; + } + d[j] = d[j] * o_scale[j] + d_local; + } +} + +template +__device__ __forceinline__ void write_p_rmem_smem(typename KTraits::SharedStorage* smem_storage, + const uint32_t stage_idx, uint32_t* p_frag) { + static constexpr uint32_t NUM_MMA_KV = KTraits::NUM_MMA_KV; + const uint32_t lane_idx = cutlass::canonical_lane_idx(); + const uint32_t warp_idx_in_wg = cutlass::canonical_warp_idx() % 4; + smem_t p_smem(smem_storage->kv_o_smem[stage_idx].p); +#pragma unroll + for (uint32_t mma_kv = 0; mma_kv < NUM_MMA_KV; ++mma_kv) { + uint32_t p_smem_offset_w = + get_swizzle_offset( + warp_idx_in_wg * 16 + lane_idx % 16, mma_kv * 2 + lane_idx / 16); + p_smem.stmatrix_m8n8x4(p_smem_offset_w, p_frag + mma_kv * 4); + } +} + +template +__device__ __forceinline__ void normalize_d_(typename KTraits::SharedStorage* smem_storage, + float* o_frag, float* m, float* d) { + float d_rcp[2]; + // compute reciprocal of d +#pragma unroll + for (uint32_t j = 0; j < 2; ++j) { + d_rcp[j] = (m[j] != -math::inf) ? math::ptx_rcp(d[j]) : 0.f; + } + +#pragma unroll + for (uint32_t reg_id = 0; reg_id < KTraits::NUM_REGS_O_FRAG; ++reg_id) { + o_frag[reg_id] = o_frag[reg_id] * d_rcp[(reg_id % 4) / 2]; + } +} + +template +__device__ __forceinline__ void write_o(typename KTraits::SharedStorage* smem_storage, + const uint32_t stage_idx, typename KTraits::DTypeO* final_o, + float* final_lse, typename KTraits::DTypeO* partial_o, + float* partial_lse, float(*o_frag), float* m, float* d, + const uint32_t o_stride_n, const uint32_t o_stride_h, + const uint32_t q_len, const uint32_t packed_offset, + const uint_fastdiv& num_heads) { + using DTypeO = typename KTraits::DTypeO; + constexpr uint32_t NUM_MMA_D_CKV = KTraits::NUM_MMA_D_CKV; + constexpr uint32_t HEAD_DIM_CKV = KTraits::HEAD_DIM_CKV; + constexpr uint32_t UPCAST_STRIDE_FINAL_O = KTraits::UPCAST_STRIDE_FINAL_O; + const uint32_t lane_idx = cutlass::canonical_lane_idx(); + const uint32_t warp_group_idx = cutlass::canonical_warp_group_idx(); + const uint32_t warp_idx_in_wg = cutlass::canonical_warp_idx() % 4; + smem_t o_smem; + o_smem = smem_storage->kv_o_smem[stage_idx].o; + + // step 0. rmem to smem +#pragma unroll + for (uint32_t k = 0; k < HEAD_DIM_CKV / 32; ++k) { + uint32_t o_frag_f16[8 / 2]; + vec_cast::cast<8>((DTypeO*)o_frag_f16, &o_frag[k * 8]); + uint32_t o_smem_offset_w = get_swizzle_offset( + warp_idx_in_wg * 16 + lane_idx % 16, + warp_group_idx * NUM_MMA_D_CKV + k * 2 + lane_idx / 16); + o_smem.template stmatrix_m8n8x4(o_smem_offset_w, o_frag_f16); + } + + if (partial_o != nullptr) { + // NOTE(Zihao): o_smem is not used if write to partial_o, and we can avoid the barrier + // write to partial_o + +#pragma unroll + for (uint32_t j = 0; j < 4; ++j) { + uint32_t q_idx = (packed_offset + warp_idx_in_wg * 16 + 4 * j + lane_idx / 8) / num_heads; + DTypeO* o_partial_ptr = + partial_o + + ((blockIdx.x * 4 + warp_idx_in_wg) * 16 + 4 * j + lane_idx / 8) * HEAD_DIM_CKV + + warp_group_idx * (HEAD_DIM_CKV / 2) + (lane_idx % 8) * upcast_size(); + uint32_t o_smem_offset_w = get_swizzle_offset( + warp_idx_in_wg * 16 + 4 * j + lane_idx / 8, + warp_group_idx * NUM_MMA_D_CKV + lane_idx % 8); +#pragma unroll + for (uint32_t k = 0; k < HEAD_DIM_CKV / 128; ++k) { + if (q_idx < q_len) { + o_smem.template store_128b(o_smem_offset_w, o_partial_ptr); + } + o_partial_ptr += 8 * upcast_size(); + o_smem_offset_w += 64; + } + } + + if constexpr (write_lse) { +#pragma unroll + for (uint32_t j = 0; j < 2; ++j) { + uint32_t q_idx = (packed_offset + warp_idx_in_wg * 16 + 8 * j + lane_idx / 4) / num_heads; + if (lane_idx % 4 == 0 && q_idx < q_len) { + partial_lse[(blockIdx.x * 4 + warp_idx_in_wg) * 16 + 8 * j + lane_idx / 4] = + math::ptx_log2(d[j]) + float(m[j]); + } + } + } + } else { + // write to final_o + +// step 1. smem to gmem +#pragma unroll + for (uint32_t j = 0; j < 4; ++j) { + uint32_t q, r; + num_heads.divmod(packed_offset + warp_idx_in_wg * 16 + 4 * j + lane_idx / 8, q, r); + DTypeO* o_final_ptr = final_o + q * o_stride_n + r * o_stride_h + + warp_group_idx * (HEAD_DIM_CKV / 2) + + (lane_idx % 8) * upcast_size(); + uint32_t o_smem_offset_w = get_swizzle_offset( + warp_idx_in_wg * 16 + 4 * j + lane_idx / 8, + warp_group_idx * NUM_MMA_D_CKV + lane_idx % 8); +#pragma unroll + for (uint32_t k = 0; k < HEAD_DIM_CKV / 128; ++k) { + if (q < q_len) { + o_smem.template store_128b(o_smem_offset_w, o_final_ptr); + } + o_final_ptr += 8 * upcast_size(); + o_smem_offset_w += 64; + } + } + + if constexpr (write_lse) { + if (final_lse) { +#pragma unroll + for (uint32_t j = 0; j < 2; ++j) { + uint32_t q, r; + num_heads.divmod(packed_offset + warp_idx_in_wg * 16 + 8 * j + lane_idx / 4, q, r); + if (lane_idx % 4 == 0 && q < q_len) { + final_lse[q * num_heads + r] = math::ptx_log2(d[j]) + float(m[j]); + } + } + } + } + } +} + +template +__device__ __forceinline__ auto get_block_coord(const Params& params, const uint32_t work_idx) { + return std::tuple(params.q_indptr[work_idx], params.kv_indptr[work_idx], + params.partial_indptr[work_idx], params.q_len[work_idx], + params.kv_len[work_idx], params.q_start[work_idx], params.kv_start[work_idx], + params.kv_end[work_idx]); +} + +template +__device__ __forceinline__ void convert_s_to_p(float* s_frag, uint32_t* p_frag) { +#pragma unroll + for (uint32_t i = 0; i < KTraits::NUM_REGS_S_FRAG / 8; ++i) { + vec_cast::cast<8>( + ((typename KTraits::DTypeKV*)p_frag) + i * 8, s_frag + i * 8); + } +} + +template +__device__ __forceinline__ void write_o_scale_smem(typename KTraits::SharedStorage* smem_storage, + float* o_scale) { + const uint32_t lane_idx = cutlass::canonical_lane_idx(); + const uint32_t warp_idx_in_wg = cutlass::canonical_warp_idx() % 4; +#pragma unroll + for (uint32_t j = 0; j < 2; ++j) { + if (lane_idx % 4 == 0) { + smem_storage->o_scale[warp_idx_in_wg * 16 + j * 8 + lane_idx / 4] = o_scale[j]; + } + } +} + +template +__device__ __forceinline__ void load_o_scale_smem(typename KTraits::SharedStorage* smem_storage, + float* o_scale) { + const uint32_t lane_idx = cutlass::canonical_lane_idx(); + const uint32_t warp_idx_in_wg = cutlass::canonical_warp_idx() % 4; +#pragma unroll + for (uint32_t j = 0; j < 2; ++j) { + o_scale[j] = smem_storage->o_scale[warp_idx_in_wg * 16 + j * 8 + lane_idx / 4]; + } +} + +template +__global__ __launch_bounds__(KTraits::NUM_THREADS) void BatchMLAPageAttentionHopperKernel( + const __grid_constant__ Params params) { + using DTypeQ = typename Params::DTypeQ; + using DTypeKV = typename Params::DTypeKV; + using DTypeO = typename Params::DTypeO; + using IdType = typename Params::IdType; + + extern __shared__ __align__(alignof(typename KTraits::SharedStorage)) uint8_t smem[]; + auto& smem_storage = reinterpret_cast(smem); + + typename KTraits::AttentionVariant variant(params, blockIdx.y, smem); + [[maybe_unused]] constexpr SwizzleMode SWIZZLE_MODE_Q_NOPE = KTraits::SWIZZLE_MODE_Q_NOPE; + [[maybe_unused]] constexpr SwizzleMode SWIZZLE_MODE_Q_PE = KTraits::SWIZZLE_MODE_Q_PE; + [[maybe_unused]] constexpr SwizzleMode SWIZZLE_MODE_CKV = KTraits::SWIZZLE_MODE_CKV; + [[maybe_unused]] constexpr SwizzleMode SWIZZLE_MODE_KPE = KTraits::SWIZZLE_MODE_KPE; + [[maybe_unused]] constexpr uint32_t NUM_MMA_KV = KTraits::NUM_MMA_KV; + [[maybe_unused]] constexpr uint32_t NUM_MMA_D_CKV = KTraits::NUM_MMA_D_CKV; + [[maybe_unused]] constexpr uint32_t HEAD_DIM_CKV = KTraits::HEAD_DIM_CKV; + [[maybe_unused]] constexpr uint32_t CTA_TILE_Q = KTraits::CTA_TILE_Q; + [[maybe_unused]] constexpr uint32_t CTA_TILE_KV = KTraits::CTA_TILE_KV; + [[maybe_unused]] constexpr int32_t NUM_STAGES = KTraits::NUM_STAGES; + [[maybe_unused]] constexpr uint32_t NUM_COPY_THREADS = KTraits::NUM_COPY_THREADS; + [[maybe_unused]] constexpr bool CAUSAL = KTraits::CAUSAL; + + DTypeQ* q_nope = params.q_nope; + DTypeQ* q_pe = params.q_pe; + DTypeKV* ckv = params.ckv; + DTypeKV* kpe = params.kpe; + IdType* kv_indices = params.kv_indices; + DTypeO* partial_o = params.partial_o; + float* partial_lse = params.partial_lse; + DTypeO* final_o = params.final_o; + float* final_lse = params.final_lse; + IdType* work_indptr = params.work_indptr; + + const uint_fastdiv& num_heads = params.num_heads; + const uint_fastdiv& block_size = params.block_size; + const uint32_t q_nope_stride_n = params.q_nope_stride_n; + const uint32_t q_nope_stride_h = params.q_nope_stride_h; + const uint32_t q_pe_stride_n = params.q_pe_stride_n; + const uint32_t q_pe_stride_h = params.q_pe_stride_h; + const uint32_t ckv_stride_page = params.ckv_stride_page; + const uint32_t ckv_stride_n = params.ckv_stride_n; + const uint32_t kpe_stride_page = params.kpe_stride_page; + const uint32_t kpe_stride_n = params.kpe_stride_n; + const uint32_t o_stride_n = params.o_stride_n; + const uint32_t o_stride_h = params.o_stride_h; + const uint32_t cluster_tile_q = gridDim.x * KTraits::CTA_TILE_Q; + + const uint32_t lane_predicate = cute::elect_one_sync(); + const uint32_t lane_idx = cutlass::canonical_lane_idx(); + const uint32_t warp_group_idx = cutlass::canonical_warp_group_idx(); + const uint32_t warp_idx = cutlass::canonical_warp_idx(); + const uint32_t warp_idx_in_wg = cutlass::canonical_warp_idx() % 4; + + PROFILER_INIT(params, smem_storage, variant, warp_group_idx, 2, (threadIdx.x % 128 == 0)); + + using MainloopPipeline = typename KTraits::MainloopPipeline; + using PipelineParams = typename MainloopPipeline::Params; + using PipelineState = typename MainloopPipeline::PipelineState; + PipelineParams pipeline_params; + pipeline_params.role = warp_group_idx == 0 ? MainloopPipeline::ThreadCategory::Producer + : MainloopPipeline::ThreadCategory::Consumer; + pipeline_params.producer_arv_count = 128; + pipeline_params.consumer_arv_count = 128; + MainloopPipeline pipeline_q(smem_storage.pipeline_q, pipeline_params); + pipeline_params.role = warp_group_idx == 0 ? MainloopPipeline::ThreadCategory::ProducerConsumer + : MainloopPipeline::ThreadCategory::Consumer; + pipeline_params.producer_arv_count = 128; + pipeline_params.consumer_arv_count = 256; + MainloopPipeline pipeline_kv(smem_storage.pipeline_kv, pipeline_params); + + __syncthreads(); + alignas(16) float o_frag[KTraits::NUM_REGS_O_FRAG]; + float m[2]; + float d[2]; + float o_scale[2]; + auto consumer_wait = [](auto& pipeline, auto& smem_pipe_read) { + auto barrier_token = pipeline.consumer_try_wait(smem_pipe_read); + pipeline.consumer_wait(smem_pipe_read, barrier_token); + }; + + if (warp_group_idx == 0) { + // load q & kv, compute pv1 + PipelineState smem_pipe_write_q = cutlass::make_producer_start_state(); + PipelineState smem_pipe_write_kv = cutlass::make_producer_start_state(); + PipelineState smem_pipe_read_kv; + + int64_t ckv_offset[KTraits::NUM_MMA_KV / 2][2]; + int64_t kpe_offset[KTraits::NUM_MMA_KV / 2][2]; + +#pragma unroll 1 + for (IdType work_idx = work_indptr[blockIdx.y]; work_idx < work_indptr[blockIdx.y + 1]; + ++work_idx) { + auto [q_indptr, kv_indptr, partial_indptr, q_len, kv_len, packed_qo_start, kv_start, kv_end] = + get_block_coord(params, work_idx); + + init_states_(o_frag, m, d, o_scale); + + const uint32_t qo_packed_idx_base = packed_qo_start + blockIdx.x * KTraits::CTA_TILE_Q; + const uint32_t qo_upperbound = + min(q_len, ceil_div(qo_packed_idx_base + KTraits::CTA_TILE_Q, num_heads)); + + uint32_t packed_kv_bound = kv_indptr * block_size + kv_len; + int kv_tile_idx = + ceil_div( + (CAUSAL ? min(kv_end, kv_len - q_len + (packed_qo_start + cluster_tile_q) / num_heads) + : kv_end), + CTA_TILE_KV) - + 1 - (kv_start / CTA_TILE_KV); + + bool has_kv = kv_tile_idx >= 0; + + const uint32_t block_iter_base = kv_indptr * block_size + kv_start; + + prefetch_offset(block_iter_base + kv_tile_idx * CTA_TILE_KV, packed_kv_bound, + ckv_stride_page, ckv_stride_n, kpe_stride_page, kpe_stride_n, + block_size, kv_indices, ckv_offset, kpe_offset); + if (has_kv) { + pipeline_kv.producer_acquire(smem_pipe_write_kv); + PROFILER_EVENT_START(variant, ProfileEventType::kIssueLoadKV); + load_kv(&smem_storage, ckv, kpe, packed_kv_bound, + block_iter_base + kv_tile_idx * CTA_TILE_KV, + smem_pipe_write_kv.index(), ckv_offset, kpe_offset); + PROFILER_EVENT_END(variant, ProfileEventType::kIssueLoadKV); + pipeline_kv.producer_commit(smem_pipe_write_kv, cutlass::arch::cpasync_barrier_arrive); + kv_tile_idx -= 1; + ++smem_pipe_write_kv; + prefetch_offset(block_iter_base + kv_tile_idx * CTA_TILE_KV, packed_kv_bound, + ckv_stride_page, ckv_stride_n, kpe_stride_page, kpe_stride_n, + block_size, kv_indices, ckv_offset, kpe_offset); + } + + pipeline_q.producer_acquire(smem_pipe_write_q); + PROFILER_EVENT_START(variant, ProfileEventType::kIssueLoadQ); + load_q(&smem_storage, q_nope + q_indptr * q_nope_stride_n, + q_pe + q_indptr * q_pe_stride_n, q_nope_stride_n, q_nope_stride_h, + q_pe_stride_n, q_pe_stride_h, qo_upperbound, qo_packed_idx_base, + params.num_heads); + PROFILER_EVENT_END(variant, ProfileEventType::kIssueLoadQ); + pipeline_q.producer_commit(smem_pipe_write_q, cutlass::arch::cpasync_barrier_arrive); + ++smem_pipe_write_q; + +#pragma unroll 1 + for (; kv_tile_idx >= 0; --kv_tile_idx) { + pipeline_kv.producer_acquire(smem_pipe_write_kv); + PROFILER_EVENT_START(variant, ProfileEventType::kIssueLoadKV); + load_kv(&smem_storage, ckv, kpe, packed_kv_bound, + block_iter_base + kv_tile_idx * CTA_TILE_KV, + smem_pipe_write_kv.index(), ckv_offset, kpe_offset); + PROFILER_EVENT_END(variant, ProfileEventType::kIssueLoadKV); + if (kv_tile_idx > 0) { + prefetch_offset(block_iter_base + (kv_tile_idx - 1) * CTA_TILE_KV, + packed_kv_bound, ckv_stride_page, ckv_stride_n, kpe_stride_page, + kpe_stride_n, block_size, kv_indices, ckv_offset, kpe_offset); + } + pipeline_kv.producer_commit(smem_pipe_write_kv, cutlass::arch::cpasync_barrier_arrive); + ++smem_pipe_write_kv; + + barrier_sync(KTraits::NUM_THREADS, NamedBarriers::kOScaleReady); + load_o_scale_smem(&smem_storage, o_scale); + PROFILER_EVENT_START(variant, ProfileEventType::kRescaleO); + rescale_o_(o_scale, o_frag); + PROFILER_EVENT_END(variant, ProfileEventType::kRescaleO); + consumer_wait(pipeline_kv, smem_pipe_read_kv); + __syncthreads(); + PROFILER_EVENT_START(variant, ProfileEventType::kGemmPV); + compute_mla_pv(&smem_storage, smem_pipe_read_kv.index(), o_frag); + warpgroup_wait<0>(); + PROFILER_EVENT_END(variant, ProfileEventType::kGemmPV); + pipeline_kv.consumer_release(smem_pipe_read_kv); + ++smem_pipe_read_kv; + } + + if (has_kv) { + barrier_sync(KTraits::NUM_THREADS, NamedBarriers::kOScaleReady); + load_o_scale_smem(&smem_storage, o_scale); + PROFILER_EVENT_START(variant, ProfileEventType::kRescaleO); + rescale_o_(o_scale, o_frag); + PROFILER_EVENT_END(variant, ProfileEventType::kRescaleO); + consumer_wait(pipeline_kv, smem_pipe_read_kv); + __syncthreads(); + PROFILER_EVENT_START(variant, ProfileEventType::kGemmPV); + compute_mla_pv(&smem_storage, smem_pipe_read_kv.index(), o_frag); + warpgroup_wait<0>(); + PROFILER_EVENT_END(variant, ProfileEventType::kGemmPV); + pipeline_kv.consumer_release(smem_pipe_read_kv); + ++smem_pipe_read_kv; + } + + barrier_sync(KTraits::NUM_THREADS, NamedBarriers::kMDReady); +#pragma unroll + for (uint32_t j = 0; j < 2; ++j) { + m[j] = smem_storage.m[warp_idx_in_wg * 16 + j * 8 + lane_idx / 4]; + d[j] = smem_storage.d[warp_idx_in_wg * 16 + j * 8 + lane_idx / 4]; + } + normalize_d_(&smem_storage, o_frag, m, d); + finalize_m_(variant, m); + PROFILER_EVENT_START(variant, ProfileEventType::kWriteO); + write_o( + &smem_storage, smem_pipe_write_kv.index(), final_o + q_indptr * o_stride_n, + final_lse ? final_lse + q_indptr * num_heads : nullptr, + (partial_indptr == -1) ? nullptr : partial_o + partial_indptr * KTraits::HEAD_DIM_CKV, + (partial_indptr == -1) ? nullptr : partial_lse + partial_indptr, o_frag, m, d, o_stride_n, + o_stride_h, qo_upperbound, qo_packed_idx_base, num_heads); + PROFILER_EVENT_END(variant, ProfileEventType::kWriteO); + __syncthreads(); + } + } else { + // compute qk, pv2 + PipelineState smem_pipe_read_q; + PipelineState smem_pipe_read_kv; + float s_frag[KTraits::NUM_REGS_S_FRAG]; + uint32_t p_frag[KTraits::NUM_REGS_P_FRAG]; + +#pragma unroll 1 + for (IdType work_idx = work_indptr[blockIdx.y]; work_idx < work_indptr[blockIdx.y + 1]; + ++work_idx) { + auto [q_indptr, kv_indptr, partial_indptr, q_len, kv_len, packed_qo_start, kv_start, kv_end] = + get_block_coord(params, work_idx); + const uint32_t qo_packed_idx_base = packed_qo_start + blockIdx.x * KTraits::CTA_TILE_Q; + const uint32_t qo_upperbound = + min(q_len, ceil_div(qo_packed_idx_base + KTraits::CTA_TILE_Q, num_heads)); + + init_states_(o_frag, m, d, o_scale); + + int kv_tile_idx = + ceil_div( + (CAUSAL ? min(kv_end, kv_len - q_len + (packed_qo_start + cluster_tile_q) / num_heads) + : kv_end), + CTA_TILE_KV) - + 1 - (kv_start / CTA_TILE_KV); + + int mask_tile_idx = + (CAUSAL ? min(kv_end, kv_len - q_len + packed_qo_start / num_heads) : kv_end) / + CTA_TILE_KV - + (kv_start / CTA_TILE_KV); + + consumer_wait(pipeline_q, smem_pipe_read_q); +#pragma unroll 1 + for (; kv_tile_idx >= mask_tile_idx && kv_tile_idx > 0; --kv_tile_idx) { + consumer_wait(pipeline_kv, smem_pipe_read_kv); + PROFILER_EVENT_START(variant, ProfileEventType::kGemmQK); + compute_mla_qk(&smem_storage, smem_pipe_read_kv.index(), s_frag); + warpgroup_wait<0>(); + PROFILER_EVENT_END(variant, ProfileEventType::kGemmQK); + logits_mask_(qo_packed_idx_base, kv_start + kv_tile_idx * CTA_TILE_KV, q_len, + kv_len, kv_end, num_heads, s_frag); + PROFILER_EVENT_START(variant, ProfileEventType::kSoftmaxUpdate); + update_md_(&smem_storage, variant, s_frag, m, d, o_scale); + PROFILER_EVENT_END(variant, ProfileEventType::kSoftmaxUpdate); + write_o_scale_smem(&smem_storage, o_scale); + + convert_s_to_p(s_frag, p_frag); + write_p_rmem_smem(&smem_storage, smem_pipe_read_kv.index(), p_frag); + barrier_arrive(KTraits::NUM_THREADS, NamedBarriers::kOScaleReady); + PROFILER_EVENT_START(variant, ProfileEventType::kRescaleO); + rescale_o_(o_scale, o_frag); + PROFILER_EVENT_END(variant, ProfileEventType::kRescaleO); + __syncthreads(); + PROFILER_EVENT_START(variant, ProfileEventType::kGemmPV); + compute_mla_pv(&smem_storage, smem_pipe_read_kv.index(), o_frag); + warpgroup_wait<0>(); + PROFILER_EVENT_END(variant, ProfileEventType::kGemmPV); + pipeline_kv.consumer_release(smem_pipe_read_kv); + ++smem_pipe_read_kv; + } + +#pragma unroll 1 + for (; kv_tile_idx + 1 > NUM_STAGES; --kv_tile_idx) { + consumer_wait(pipeline_kv, smem_pipe_read_kv); + PROFILER_EVENT_START(variant, ProfileEventType::kGemmQK); + compute_mla_qk(&smem_storage, smem_pipe_read_kv.index(), s_frag); + warpgroup_wait<0>(); + PROFILER_EVENT_END(variant, ProfileEventType::kGemmQK); + PROFILER_EVENT_START(variant, ProfileEventType::kSoftmaxUpdate); + update_md_(&smem_storage, variant, s_frag, m, d, o_scale); + PROFILER_EVENT_END(variant, ProfileEventType::kSoftmaxUpdate); + write_o_scale_smem(&smem_storage, o_scale); + convert_s_to_p(s_frag, p_frag); + write_p_rmem_smem(&smem_storage, smem_pipe_read_kv.index(), p_frag); + barrier_arrive(KTraits::NUM_THREADS, NamedBarriers::kOScaleReady); + PROFILER_EVENT_START(variant, ProfileEventType::kRescaleO); + rescale_o_(o_scale, o_frag); + PROFILER_EVENT_END(variant, ProfileEventType::kRescaleO); + __syncthreads(); + PROFILER_EVENT_START(variant, ProfileEventType::kGemmPV); + compute_mla_pv(&smem_storage, smem_pipe_read_kv.index(), o_frag); + warpgroup_wait<0>(); + PROFILER_EVENT_END(variant, ProfileEventType::kGemmPV); + pipeline_kv.consumer_release(smem_pipe_read_kv); + ++smem_pipe_read_kv; + } + +#pragma unroll 1 + for (; kv_tile_idx >= 0; --kv_tile_idx) { + consumer_wait(pipeline_kv, smem_pipe_read_kv); + PROFILER_EVENT_START(variant, ProfileEventType::kGemmQK); + compute_mla_qk(&smem_storage, smem_pipe_read_kv.index(), s_frag); + warpgroup_wait<0>(); + PROFILER_EVENT_END(variant, ProfileEventType::kGemmQK); + logits_mask_(qo_packed_idx_base, kv_start + kv_tile_idx * CTA_TILE_KV, q_len, + kv_len, kv_end, num_heads, s_frag); + PROFILER_EVENT_START(variant, ProfileEventType::kSoftmaxUpdate); + update_md_(&smem_storage, variant, s_frag, m, d, o_scale); + PROFILER_EVENT_END(variant, ProfileEventType::kSoftmaxUpdate); + write_o_scale_smem(&smem_storage, o_scale); + convert_s_to_p(s_frag, p_frag); + write_p_rmem_smem(&smem_storage, smem_pipe_read_kv.index(), p_frag); + barrier_arrive(KTraits::NUM_THREADS, NamedBarriers::kOScaleReady); + PROFILER_EVENT_START(variant, ProfileEventType::kRescaleO); + rescale_o_(o_scale, o_frag); + PROFILER_EVENT_END(variant, ProfileEventType::kRescaleO); + __syncthreads(); + PROFILER_EVENT_START(variant, ProfileEventType::kGemmPV); + compute_mla_pv(&smem_storage, smem_pipe_read_kv.index(), o_frag); + warpgroup_wait<0>(); + PROFILER_EVENT_END(variant, ProfileEventType::kGemmPV); + pipeline_kv.consumer_release(smem_pipe_read_kv); + ++smem_pipe_read_kv; + } + + pipeline_q.consumer_release(smem_pipe_read_q); + ++smem_pipe_read_q; + +#pragma unroll + for (uint32_t j = 0; j < 2; ++j) { + d[j] += __shfl_xor_sync(0x11111111, d[j], 0x2); + d[j] += __shfl_xor_sync(0x11111111, d[j], 0x1); + if (lane_idx % 4 == 0) { + smem_storage.m[warp_idx_in_wg * 16 + j * 8 + lane_idx / 4] = m[j]; + smem_storage.d[warp_idx_in_wg * 16 + j * 8 + lane_idx / 4] = d[j]; + } + } + normalize_d_(&smem_storage, o_frag, m, d); + finalize_m_(variant, m); + barrier_arrive(KTraits::NUM_THREADS, NamedBarriers::kMDReady); + PROFILER_EVENT_START(variant, ProfileEventType::kWriteO); + write_o( + &smem_storage, smem_pipe_read_kv.index(), final_o + q_indptr * o_stride_n, + final_lse ? final_lse + q_indptr * num_heads : nullptr, + (partial_indptr == -1) ? nullptr : partial_o + partial_indptr * KTraits::HEAD_DIM_CKV, + (partial_indptr == -1) ? nullptr : partial_lse + partial_indptr, o_frag, m, d, o_stride_n, + o_stride_h, qo_upperbound, qo_packed_idx_base, num_heads); + PROFILER_EVENT_END(variant, ProfileEventType::kWriteO); + __syncthreads(); + } + } + + auto grid = cg::this_grid(); + grid.sync(); + + PROFILER_EVENT_START(variant, ProfileEventType::kSplitK); + + __syncthreads(); + // the second stage, merge partial outputs + DevicePersistentMergeStates( + params.merge_packed_offset_start, params.merge_packed_offset_end, + params.merge_partial_packed_offset_start, params.merge_partial_packed_offset_end, + params.merge_partial_stride, partial_o, partial_lse, final_o, final_lse, o_stride_n, + o_stride_h, num_heads); + + PROFILER_EVENT_END(variant, ProfileEventType::kSplitK); +} + +} // namespace hopper + +template +cudaError_t BatchMLAPageAttentionHopper(Params params, uint32_t num_blks_x, uint32_t num_blks_y, + cudaStream_t stream) { + using DTypeQ = typename Params::DTypeQ; + using DTypeKV = typename Params::DTypeKV; + using DTypeO = typename Params::DTypeO; + using IdType = typename Params::IdType; + + if (MASK_MODE == MaskMode::kCustom) { + return cudaErrorNotSupported; + } + constexpr bool CAUSAL = MASK_MODE == MaskMode::kCausal; + + // get GPU shared memory size + int device; + int smem_limit_per_sm; + cudaGetDevice(&device); + cudaDeviceGetAttribute(&smem_limit_per_sm, cudaDevAttrMaxSharedMemoryPerMultiprocessor, device); + + constexpr uint32_t NUM_STAGES = 2; + constexpr uint32_t CTA_TILE_Q = 64; + constexpr uint32_t CTA_TILE_KV = 64; + + using KTraits = + hopper::HopperKernelTraits; + dim3 nblks(num_blks_x, num_blks_y); + dim3 nthrs(KTraits::NUM_THREADS); + size_t smem_size = sizeof(typename KTraits::SharedStorage); + + auto kernel = hopper::BatchMLAPageAttentionHopperKernel; + void* args[] = {(void*)¶ms}; + + FLASHINFER_CUDA_CALL( + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + FLASHINFER_CUDA_CALL( + cudaLaunchCooperativeKernel((void*)kernel, nblks, nthrs, args, smem_size, stream)); + + return cudaSuccess; +} + +} // namespace mla + +} // namespace flashinfer + +#endif // FLASHINFER_MLA_HOPPER_CUH_ diff --git a/include/flashinfer/attention/mla_params.cuh b/include/flashinfer/attention/mla_params.cuh new file mode 100644 index 0000000000000000000000000000000000000000..ff5d168ba20d1e4b0267d2349793b0195cb21188 --- /dev/null +++ b/include/flashinfer/attention/mla_params.cuh @@ -0,0 +1,78 @@ +/* + * Copyright (c) 2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef FLASHINFER_MLA_PARAMS_CUH_ +#define FLASHINFER_MLA_PARAMS_CUH_ +#include + +#include "../fastdiv.cuh" +#include "../profiler.cuh" + +namespace flashinfer { + +template +struct MLAParams { + using DTypeQ = DTypeQ_; + using DTypeKV = DTypeKV_; + using DTypeO = DTypeO_; + using IdType = IdType_; + + DTypeQ* q_nope; + DTypeQ* q_pe; + DTypeKV* ckv; + DTypeKV* kpe; + DTypeO* partial_o; + float* partial_lse; + DTypeO* final_o; + float* final_lse; + + IdType* q_indptr; + IdType* kv_indptr; + IdType* partial_indptr; + IdType* merge_packed_offset_start; + IdType* merge_packed_offset_end; + IdType* merge_partial_packed_offset_start; + IdType* merge_partial_packed_offset_end; + IdType* merge_partial_stride; + IdType* kv_indices; + IdType* q_len; + IdType* kv_len; + IdType* q_start; + IdType* kv_start; + IdType* kv_end; + IdType* work_indptr; + + PROFILER_PARAMS_DECL + + uint_fastdiv block_size; + uint_fastdiv num_heads; + + uint32_t q_nope_stride_n; + uint32_t q_nope_stride_h; + uint32_t q_pe_stride_n; + uint32_t q_pe_stride_h; + uint32_t ckv_stride_page; + uint32_t ckv_stride_n; + uint32_t kpe_stride_page; + uint32_t kpe_stride_n; + uint32_t o_stride_n; + uint32_t o_stride_h; + + float sm_scale; +}; + +}; // namespace flashinfer + +#endif // FLASHINFER_MLA_PARAMS_CUH_ diff --git a/include/flashinfer/attention/persistent.cuh b/include/flashinfer/attention/persistent.cuh new file mode 100644 index 0000000000000000000000000000000000000000..68a692960f94b4836fe38a7f2c102123f392197a --- /dev/null +++ b/include/flashinfer/attention/persistent.cuh @@ -0,0 +1,612 @@ +/* + * Copyright (c) 2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef FLASHINFER_PERSISTENT_CUH_ +#define FLASHINFER_PERSISTENT_CUH_ + +#include "../cp_async.cuh" +#include "../math.cuh" +#include "../utils.cuh" +#include "mask.cuh" +#include "persistent_template.cuh" +#include "prefill.cuh" +#include "state.cuh" + +namespace flashinfer { + +using cp_async::PrefetchMode; +using cp_async::SharedMemFillMode; + +template +__device__ __forceinline__ auto get_block_coord(const Params& params, const uint32_t work_idx) { + return std::tuple(params.q_indptr[work_idx], params.kv_indptr[work_idx], + params.partial_indptr[work_idx], params.q_len[work_idx], + params.kv_len[work_idx], params.q_start[work_idx], params.kv_start[work_idx], + params.kv_end[work_idx], params.kv_head_idx_arr[work_idx], + params.len_kv_chunk[work_idx]); +} + +template +__device__ __forceinline__ void prefetch_offest( + const uint32_t packed_block_iter_base, const uint32_t packed_kv_bound, + const uint32_t kv_head_idx, const uint32_t kv_stride_page, const uint32_t kv_stride_h, + const uint32_t kv_stride_n, const uint_fastdiv& block_size, typename KTraits::IdType* indices, + size_t* kv_offset) { + using DTypeKV = typename KTraits::DTypeKV; + constexpr uint32_t KV_THR_LAYOUT_ROW = KTraits::KV_THR_LAYOUT_ROW; + constexpr uint32_t KV_THR_LAYOUT_COL = KTraits::KV_THR_LAYOUT_COL; + constexpr uint32_t NUM_WARPS_Q = KTraits::NUM_WARPS_Q; + constexpr uint32_t NUM_WARPS_KV = KTraits::NUM_WARPS_KV; + constexpr uint32_t NUM_MMA_KV = KTraits::NUM_MMA_KV; + constexpr SwizzleMode SWIZZLE_MODE_KV = KTraits::SWIZZLE_MODE_KV; + const uint32_t lane_idx = threadIdx.x % 32, warp_idx = threadIdx.x / 32; + +#pragma unroll + for (uint32_t i = 0; + i < NUM_MMA_KV * (SWIZZLE_MODE_KV == SwizzleMode::k128B ? 4 : 2) / NUM_WARPS_Q; ++i) { + uint32_t page_iter, entry_idx; + uint32_t packed_block_iter = packed_block_iter_base + warp_idx * KV_THR_LAYOUT_ROW + + lane_idx / KV_THR_LAYOUT_COL + + KV_THR_LAYOUT_ROW * NUM_WARPS_Q * NUM_WARPS_KV * i; + block_size.divmod(packed_block_iter, page_iter, entry_idx); + kv_offset[i] = (packed_block_iter < packed_kv_bound ? indices[page_iter] : 0) * kv_stride_page + + entry_idx * kv_stride_n + kv_head_idx * kv_stride_h + + (lane_idx % KV_THR_LAYOUT_COL) * upcast_size(); + } +} + +template +__device__ __forceinline__ void write_o_(float (*o_frag)[KTraits::NUM_MMA_D_VO][8], + smem_t* o_smem, + typename KTraits::DTypeO* o_ptr_base, + const uint32_t o_packed_idx_base_warp, + const uint32_t o_packed_idx_base_cta, + const uint32_t qo_upper_bound, const uint32_t o_stride_n, + const uint_fastdiv group_size, const uint32_t warp_idx, + const uint32_t lane_idx, const dim3 tid) { + using DTypeO = typename KTraits::DTypeO; + constexpr uint32_t UPCAST_STRIDE_O = KTraits::UPCAST_STRIDE_O; + const uint32_t warp_idx_x = get_warp_idx_q(tid.y), + warp_idx_z = get_warp_idx_kv(tid.z); + + static_assert(sizeof(DTypeO) == 2); + if (warp_idx_z == 0) { +#pragma unroll + for (uint32_t mma_q = 0; mma_q < KTraits::NUM_MMA_Q; ++mma_q) { +#pragma unroll + for (uint32_t mma_d = 0; mma_d < KTraits::NUM_MMA_D_VO; ++mma_d) { + uint32_t o_frag_f16[8 / 2]; + vec_cast::cast<8>((DTypeO*)o_frag_f16, o_frag[mma_q][mma_d]); + +#ifdef FLASHINFER_STMATRIX_M8N8X4_ENABLED + uint32_t o_smem_offset_w = o_smem->get_permuted_offset( + (warp_idx_x * KTraits::NUM_MMA_Q + mma_q) * 16 + lane_idx % 16, + mma_d * 2 + lane_idx / 16); + o_smem->stmatrix_m8n8x4(o_smem_offset_w, o_frag_f16); +#else + uint32_t o_smem_offset_w = o_smem->get_permuted_offset( + (warp_idx_x * KTraits::NUM_MMA_Q + mma_q) * 16 + lane_idx / 4, mma_d * 2); + ((uint32_t*)(o_smem->base + o_smem_offset_w))[lane_idx % 4] = o_frag_f16[0]; + ((uint32_t*)(o_smem->base + o_smem_offset_w + 8 * UPCAST_STRIDE_O))[lane_idx % 4] = + o_frag_f16[1]; + ((uint32_t*)(o_smem->base + (o_smem_offset_w ^ 0x1)))[lane_idx % 4] = o_frag_f16[2]; + ((uint32_t*)(o_smem->base + (o_smem_offset_w ^ 0x1) + 8 * UPCAST_STRIDE_O))[lane_idx % 4] = + o_frag_f16[3]; +#endif + } + } + + uint32_t o_smem_offset_w = o_smem->get_permuted_offset( + warp_idx_x * KTraits::NUM_MMA_Q * 16 + lane_idx / 8, lane_idx % 8); + +#pragma unroll + for (uint32_t mma_q = 0; mma_q < KTraits::NUM_MMA_Q; ++mma_q) { +#pragma unroll + for (uint32_t j = 0; j < 2 * 2; ++j) { + uint32_t q, r; + const uint32_t o_packed_idx = o_packed_idx_base_warp + lane_idx / 8 + mma_q * 16 + j * 4; + group_size.divmod(o_packed_idx, q, r); + + const uint32_t o_idx = q; + DTypeO* o_ptr = o_ptr_base + (o_packed_idx - o_packed_idx_base_cta) * o_stride_n + + (lane_idx % 8) * upcast_size(); +#pragma unroll + for (uint32_t mma_do = 0; mma_do < KTraits::NUM_MMA_D_VO / 4; ++mma_do) { + if (o_idx < qo_upper_bound) { + o_smem->store_128b(o_smem_offset_w, o_ptr); + } + o_ptr += 8 * upcast_size(); + o_smem_offset_w = o_smem->template advance_offset_by_column<8>(o_smem_offset_w, mma_do); + } + o_smem_offset_w = + o_smem->template advance_offset_by_row<4, UPCAST_STRIDE_O>(o_smem_offset_w) - + 2 * KTraits::NUM_MMA_D_VO; + } + } + } +} + +template +struct BlockBatchPagedAttentionPersistent { + using KTraits = KTraits_; + using Params = Params_; + + static __device__ __forceinline__ void Run(const Params& params, + typename KTraits::SharedStorage* smem_storage + PROFILER_CLOSURE_FUNC_PARAMS) { + using DTypeQ = typename Params::DTypeQ; + using DTypeKV = typename Params::DTypeKV; + using DTypeO = typename Params::DTypeO; + using IdType = typename Params::IdType; + using DTypeQKAccum = typename KTraits::DTypeQKAccum; + using AttentionVariant = typename KTraits::AttentionVariant; + [[maybe_unused]] constexpr uint32_t NUM_MMA_Q = KTraits::NUM_MMA_Q; + [[maybe_unused]] constexpr uint32_t NUM_MMA_KV = KTraits::NUM_MMA_KV; + [[maybe_unused]] constexpr uint32_t NUM_MMA_D_QK = KTraits::NUM_MMA_D_QK; + [[maybe_unused]] constexpr uint32_t NUM_MMA_D_VO = KTraits::NUM_MMA_D_VO; + [[maybe_unused]] constexpr uint32_t HEAD_DIM_QK = KTraits::HEAD_DIM_QK; + [[maybe_unused]] constexpr uint32_t HEAD_DIM_VO = KTraits::HEAD_DIM_VO; + [[maybe_unused]] constexpr uint32_t UPCAST_STRIDE_Q = KTraits::UPCAST_STRIDE_Q; + [[maybe_unused]] constexpr uint32_t UPCAST_STRIDE_K = KTraits::UPCAST_STRIDE_K; + [[maybe_unused]] constexpr uint32_t UPCAST_STRIDE_V = KTraits::UPCAST_STRIDE_V; + [[maybe_unused]] constexpr uint32_t UPCAST_STRIDE_O = KTraits::UPCAST_STRIDE_O; + [[maybe_unused]] constexpr uint32_t NUM_WARPS_Q = KTraits::NUM_WARPS_Q; + [[maybe_unused]] constexpr uint32_t NUM_WARPS_KV = KTraits::NUM_WARPS_KV; + [[maybe_unused]] constexpr SwizzleMode SWIZZLE_MODE_Q = KTraits::SWIZZLE_MODE_Q; + [[maybe_unused]] constexpr SwizzleMode SWIZZLE_MODE_KV = KTraits::SWIZZLE_MODE_KV; + [[maybe_unused]] constexpr uint32_t CTA_TILE_Q = KTraits::CTA_TILE_Q; + [[maybe_unused]] constexpr uint32_t CTA_TILE_KV = KTraits::CTA_TILE_KV; + [[maybe_unused]] constexpr bool CAUSAL = KTraits::MASK_MODE == MaskMode::kCausal; + [[maybe_unused]] constexpr uint32_t NUM_STAGES = KTraits::NUM_STAGES; + + DTypeQ* q = params.q; + DTypeKV* k = params.k; + DTypeKV* v = params.v; + IdType* kv_indices = params.kv_indices; + float* partial_lse = params.partial_lse; + IdType* work_indptr = params.work_indptr; + + float s_frag[NUM_MMA_Q][NUM_MMA_KV][8]; + alignas(16) float o_frag[NUM_MMA_Q][NUM_MMA_D_VO][8]; + float m[NUM_MMA_Q][2]; + float d[NUM_MMA_Q][2]; + + const uint_fastdiv& gqa_group_size = params.gqa_group_size; + const uint32_t num_kv_heads = params.num_kv_heads; + const uint_fastdiv& block_size = params.page_size; + const uint32_t q_stride_n = params.q_stride_n; + const uint32_t q_stride_h = params.q_stride_h; + const uint32_t k_stride_page = params.k_stride_page; + const uint32_t k_stride_h = params.k_stride_h; + const uint32_t k_stride_n = params.k_stride_n; + const uint32_t v_stride_page = params.v_stride_page; + const uint32_t v_stride_h = params.v_stride_h; + const uint32_t v_stride_n = params.v_stride_n; + const uint32_t cluster_tile_q = gridDim.x * CTA_TILE_Q; + smem_t q_smem(smem_storage->q_smem); + + AttentionVariant variant(params, /*batch_idx=*/0, nullptr); + + const uint32_t lane_idx = threadIdx.x % 32; + const uint32_t warp_idx = threadIdx.x / 32; + + // threadIdx: [32, NUM_WARPS_Q, NUM_WARPS_KV] + // remap to utilize tool function in FA2 prefill + const dim3 tid = dim3(lane_idx, warp_idx % NUM_WARPS_Q, warp_idx / NUM_WARPS_Q); + + uint32_t q_smem_offset_r = get_permuted_offset( + get_warp_idx_q(tid.y) * NUM_MMA_Q * 16 + lane_idx % 16, lane_idx / 16); + uint32_t k_smem_offset_r = get_permuted_offset( + get_warp_idx_kv(tid.z) * NUM_MMA_KV * 16 + 8 * (lane_idx / 16) + + lane_idx % 8, + (lane_idx % 16) / 8), + v_smem_offset_r = get_permuted_offset( + get_warp_idx_kv(tid.z) * NUM_MMA_KV * 16 + lane_idx % 16, lane_idx / 16); + uint32_t k_smem_offset_w = get_permuted_offset( + warp_idx * KTraits::KV_THR_LAYOUT_ROW + lane_idx / KTraits::KV_THR_LAYOUT_COL, + lane_idx % KTraits::KV_THR_LAYOUT_COL), + v_smem_offset_w = get_permuted_offset( + warp_idx * KTraits::KV_THR_LAYOUT_ROW + lane_idx / KTraits::KV_THR_LAYOUT_COL, + lane_idx % KTraits::KV_THR_LAYOUT_COL); + size_t thr_local_kv_offset[NUM_MMA_KV * KTraits::KV_THR_LAYOUT_COL / 2 / KTraits::NUM_WARPS_Q]; + +#pragma unroll 1 + for (IdType work_idx = work_indptr[blockIdx.y]; work_idx < work_indptr[blockIdx.y + 1]; + ++work_idx) { + // profile log + if constexpr (CTA_TILE_Q > 64) { + PROFILER_EVENT_START(profiler_closure, PersistentProfileEventType::kRunner1); + } else { + PROFILER_EVENT_START(profiler_closure, PersistentProfileEventType::kRunner2); + } + + const auto [q_indptr, kv_indptr, o_indptr, q_len, kv_len, packed_qo_start, kv_start, kv_end, + kv_head_idx, len_kv_chunk] = get_block_coord(params, work_idx); + + const uint32_t kv_chunk_idx = ceil_div(kv_start, len_kv_chunk); + const uint32_t num_kv_chunks = ceil_div( + CAUSAL + ? min((kv_len - q_len) + (packed_qo_start + cluster_tile_q) / gqa_group_size, kv_len) + : kv_len, + len_kv_chunk); + + const uint32_t qo_packed_idx_base = packed_qo_start + blockIdx.x * CTA_TILE_Q + + get_warp_idx_q(tid.y) * NUM_MMA_Q * 16; + const uint32_t qo_upperbound = + min(q_len, ceil_div(qo_packed_idx_base + CTA_TILE_Q, gqa_group_size)); + + init_states(variant, o_frag, m, d); + + DTypeQ* q_ptr_base = q + q_indptr * q_stride_n + (kv_head_idx * gqa_group_size) * q_stride_h; + + // load_q + load_q_global_smem(qo_packed_idx_base, qo_upperbound, q_ptr_base, q_stride_n, + q_stride_h, gqa_group_size, &q_smem, tid); + + smem_t k_smem(smem_storage->k_smem), v_smem(smem_storage->v_smem); + int kv_tile_idx = + ceil_div((CAUSAL ? min(kv_end, + kv_len - q_len + + ceil_div((packed_qo_start + cluster_tile_q), gqa_group_size)) + : kv_end), + CTA_TILE_KV) - + 1 - (kv_start / CTA_TILE_KV); + + int mask_tile_idx = + (CAUSAL ? min(kv_end, kv_len - q_len + ceil_div(packed_qo_start, gqa_group_size)) + : kv_end) / + CTA_TILE_KV - + (kv_start / CTA_TILE_KV); + + uint32_t block_iter_base = kv_indptr * block_size + kv_start; + // last kv tile + __syncthreads(); + uint32_t packed_kv_bound = kv_indptr * block_size + kv_len; + + prefetch_offest(block_iter_base + kv_tile_idx * CTA_TILE_KV, packed_kv_bound, + kv_head_idx, k_stride_page, k_stride_h, k_stride_n, block_size, + kv_indices, thr_local_kv_offset); + page_produce_kv(smem_storage, &k_smem_offset_w, k, + kv_start + kv_tile_idx * CTA_TILE_KV, thr_local_kv_offset, + kv_end, warp_idx, lane_idx); + cp_async::commit_group(); + page_produce_kv(smem_storage, &v_smem_offset_w, v, + kv_start + kv_tile_idx * CTA_TILE_KV, thr_local_kv_offset, + kv_end, warp_idx, lane_idx); + cp_async::commit_group(); + + // loop with mask + LOOP_SPLIT_MASK( + kv_tile_idx, kv_tile_idx >= mask_tile_idx && kv_tile_idx > 0, + kv_tile_idx + 1 > NUM_STAGES, { + prefetch_offest(block_iter_base + (kv_tile_idx - 1) * CTA_TILE_KV, + packed_kv_bound, kv_head_idx, k_stride_page, k_stride_h, + k_stride_n, block_size, kv_indices, thr_local_kv_offset); + cp_async::wait_group<1>(); + __syncthreads(); + + compute_qk(&q_smem, &q_smem_offset_r, &k_smem, &k_smem_offset_r, s_frag); + if constexpr (AttentionVariant::use_logits_soft_cap) { + logits_transform( + params, variant, /*batch_idx=*/0, qo_packed_idx_base, + kv_start + (kv_tile_idx * NUM_WARPS_KV + get_warp_idx_kv(tid.z)) * + NUM_MMA_KV * 16, + q_len, kv_len, gqa_group_size, s_frag, tid, kv_head_idx); + } + if constexpr (WITH_MASK) { + logits_mask( + params, variant, /*batch_idx=*/0, qo_packed_idx_base, + kv_start + (kv_tile_idx * NUM_WARPS_KV + get_warp_idx_kv(tid.z)) * + NUM_MMA_KV * 16, + q_len, kv_len, kv_end, gqa_group_size, s_frag, tid, kv_head_idx); + } + update_mdo_states(variant, s_frag, o_frag, m, d); + + __syncthreads(); + page_produce_kv(smem_storage, &k_smem_offset_w, k, + kv_start + (kv_tile_idx - 1) * CTA_TILE_KV, + thr_local_kv_offset, kv_end, warp_idx, lane_idx); + cp_async::commit_group(); + cp_async::wait_group<1>(); + + __syncthreads(); + compute_sfm_v(&v_smem, &v_smem_offset_r, s_frag, o_frag, d); + __syncthreads(); + + page_produce_kv(smem_storage, &v_smem_offset_w, v, + kv_start + (kv_tile_idx - 1) * CTA_TILE_KV, + thr_local_kv_offset, kv_end, warp_idx, lane_idx); + cp_async::commit_group(); + }); + cp_async::wait_group<0>(); + __syncthreads(); + +#pragma unroll + for (; kv_tile_idx >= 0; --kv_tile_idx) { + compute_qk(&q_smem, &q_smem_offset_r, &k_smem, &k_smem_offset_r, s_frag); + if constexpr (AttentionVariant::use_logits_soft_cap) { + logits_transform( + params, variant, /*batch_idx=*/0, qo_packed_idx_base, + kv_start + + (kv_tile_idx * NUM_WARPS_KV + get_warp_idx_kv(tid.z)) * NUM_MMA_KV * 16, + q_len, kv_len, gqa_group_size, s_frag, tid, kv_head_idx); + } + logits_mask( + params, variant, /*batch_idx=*/0, qo_packed_idx_base, + kv_start + + (kv_tile_idx * NUM_WARPS_KV + get_warp_idx_kv(tid.z)) * NUM_MMA_KV * 16, + q_len, kv_len, kv_end, gqa_group_size, s_frag, tid, kv_head_idx); + update_mdo_states(variant, s_frag, o_frag, m, d); + compute_sfm_v(&v_smem, &v_smem_offset_r, s_frag, o_frag, d); + } + + __syncthreads(); + + finalize_m(variant, m); + + // threadblock synchronization + threadblock_sync_mdo_states(o_frag, smem_storage, m, d, warp_idx, lane_idx, tid); + + // normalize d + normalize_d(o_frag, m, d); + + // write back to global memory + // o_indptr (partial_o): [packed_qo_len * num_kv_chunks, num_kv_heads, head_dim] + // q_indpt (final_o): [qo_len, num_kv_heads, gqa_group_size, head_dim] + if (num_kv_chunks > 1) { + DTypeO* o_ptr_base = params.partial_o + + ((o_indptr + kv_chunk_idx) * num_kv_heads + kv_head_idx) * HEAD_DIM_VO; + write_o_(o_frag, &q_smem, o_ptr_base, qo_packed_idx_base, packed_qo_start, + qo_upperbound, num_kv_chunks * num_kv_heads * HEAD_DIM_VO, gqa_group_size, + warp_idx, lane_idx, tid); + } else { + // write through + DTypeO* o_ptr_base = + params.final_o + q_indptr * q_stride_n + (kv_head_idx * gqa_group_size) * q_stride_h; + write_o_reg_gmem(o_frag, &q_smem, o_ptr_base, qo_packed_idx_base, q_len, + q_stride_n, q_stride_h, gqa_group_size, tid); + } + + if constexpr (variant.use_softmax) { + if (get_warp_idx_kv(tid.z) == 0) { +#pragma unroll + for (uint32_t mma_q = 0; mma_q < NUM_MMA_Q; ++mma_q) { +#pragma unroll + for (uint32_t j = 0; j < 2; ++j) { + uint32_t q, r; + const uint32_t packed_qo_idx = qo_packed_idx_base + lane_idx / 4 + j * 8 + mma_q * 16; + gqa_group_size.divmod(packed_qo_idx, q, r); + if (q < qo_upperbound) { + if (num_kv_chunks > 1) { + partial_lse[(o_indptr + (packed_qo_idx - packed_qo_start) * num_kv_chunks + + kv_chunk_idx) * + num_kv_heads + + kv_head_idx] = math::ptx_log2(d[mma_q][j]) + float(m[mma_q][j]); + } else if (params.final_lse != nullptr) { + // write through + const uint32_t qo_head_idx = kv_head_idx * gqa_group_size + r; + params.final_lse[(q_indptr + q) * num_kv_heads * gqa_group_size + qo_head_idx] = + math::ptx_log2(d[mma_q][j]) + float(m[mma_q][j]); + } + } + } + } + } + } + + // profile + if constexpr (CTA_TILE_Q > 64) { + PROFILER_EVENT_END(profiler_closure, PersistentProfileEventType::kRunner1); + } else { + PROFILER_EVENT_END(profiler_closure, PersistentProfileEventType::kRunner2); + } + } + } +}; + +template +struct StateReductionKernelTraits { + using DTypeIn = DTypeIn_; + using DTypeO = DTypeO_; + using IdType = IdType_; + + static constexpr uint32_t HEAD_DIM_VO = HEAD_DIM_VO_; + static constexpr uint32_t NUM_SMEM_STAGES = NUM_SMEM_STAGES_; + static constexpr uint32_t NUM_THREADS = NUM_THREADS_; + static constexpr uint32_t NUM_WARPS = NUM_THREADS / 32; + + static constexpr uint32_t vec_size = (16U / sizeof(DTypeIn)) > (HEAD_DIM_VO / 32U) + ? (16U / sizeof(DTypeIn)) + : (HEAD_DIM_VO / 32U); + static constexpr uint32_t bdx = HEAD_DIM_VO / vec_size; + + // gridDim is accessed by runtime variable and should be set by core attention + // workload layout [bdx, bdy, num_warps] + static_assert(NUM_THREADS % bdx == 0); + static constexpr uint32_t bdy = 32 / bdx; + + // pipeline load & reduction + static constexpr size_t SMEM_SIZE = + NUM_WARPS * NUM_SMEM_STAGES * bdy * HEAD_DIM_VO * sizeof(DTypeIn) + + NUM_THREADS * sizeof(float); +}; + +template +struct BlockBatchReductionPersistent { + using KTraits = KTraits_; + + static __device__ __forceinline__ void Run( + typename KTraits::DTypeIn* __restrict__ V, typename KTraits::DTypeO* __restrict__ v_merged, + float* __restrict__ S, float* __restrict__ s_merged, + const typename KTraits::IdType num_packed_qo_len, const uint_fastdiv gqa_group_size, + const uint32_t num_kv_heads, const typename KTraits::IdType* indptr, + const typename KTraits::IdType* o_indices, uint8_t* smem PROFILER_CLOSURE_FUNC_PARAMS) { + using DTypeIn = typename KTraits::DTypeIn; + using DTypeO = typename KTraits::DTypeO; + using IdType = typename KTraits::IdType; + + [[maybe_unused]] constexpr uint32_t bdx = KTraits::bdx; + [[maybe_unused]] constexpr uint32_t bdy = KTraits::bdy; + [[maybe_unused]] constexpr uint32_t num_warps = KTraits::NUM_WARPS; + + [[maybe_unused]] constexpr uint32_t vec_size = KTraits::vec_size; + [[maybe_unused]] constexpr uint32_t head_dim = KTraits::HEAD_DIM_VO; + [[maybe_unused]] constexpr uint32_t num_smem_stages = KTraits::NUM_SMEM_STAGES; + [[maybe_unused]] constexpr uint32_t vec_bits = sizeof(DTypeIn) * vec_size * 8; + + // control flow metadata + const uint32_t warp_idx = threadIdx.x / 32; + const uint32_t tx = (threadIdx.x % 32) % bdx, ty = (threadIdx.x % 32) / bdx; + + const uint32_t worker_id = blockIdx.y * num_warps + warp_idx; + const uint32_t num_workers = gridDim.x * gridDim.y * gridDim.z * num_warps; + + DTypeIn* v_smem = (DTypeIn*)smem + warp_idx * num_smem_stages * bdy * head_dim; + // FIXME: fix the offset calculation + float* s_smem = (float*)(smem + num_warps * num_smem_stages * bdy * head_dim * sizeof(DTypeIn) + + warp_idx * 32 * sizeof(float)); + + // V: [num_packed_qo_len x num_kv_tiles, num_kv_heads, head_dim] + // v_merged: [qo_len, num_kv_heads, gqa_group_size, head_dim] +#pragma unroll 1 + for (uint32_t i = worker_id; i < num_packed_qo_len * num_kv_heads; i += num_workers) { + PROFILER_EVENT_START(profiler_closure, PersistentProfileEventType::kReduction); + + // remap workload + uint32_t packed_qo_idx = i / num_kv_heads; + uint32_t kv_head_idx = i % num_kv_heads; + uint32_t qo_head_idx = packed_qo_idx % gqa_group_size; + + // index calculation + auto partial_idx_to_offset = [&](uint32_t off) { + return (indptr[packed_qo_idx] + off) * num_kv_heads + kv_head_idx; + }; + auto merge_idx_to_offset = [&]() { + return (o_indices[packed_qo_idx] * num_kv_heads + kv_head_idx) * gqa_group_size + + qo_head_idx; + }; + + state_t st; + const uint32_t num_index_sets = indptr[packed_qo_idx + 1] - indptr[packed_qo_idx]; + + if (num_index_sets == 0 || num_index_sets == 1) { + // already write through, bypass + PROFILER_EVENT_END(profiler_closure, PersistentProfileEventType::kReduction); + continue; + } + +#pragma unroll + for (uint32_t iter = 0; iter < num_smem_stages; ++iter) { + cp_async::pred_load( + v_smem + (iter * bdy + ty) * head_dim + tx * vec_size, + V + partial_idx_to_offset(iter * bdy + ty) * head_dim + tx * vec_size, + (iter * bdy + ty) < num_index_sets); + cp_async::commit_group(); + } +#pragma unroll 4 + for (uint32_t iter = 0; iter < ceil_div(num_index_sets, bdy); ++iter) { + if (iter % bdx == 0) { + s_smem[ty * bdx + tx] = iter * bdy + (ty * bdx + tx) < num_index_sets + ? S[partial_idx_to_offset(iter * bdy + ty * bdx + tx)] + : 0.f; + __syncwarp(); + } + cp_async::wait_group(); + __syncwarp(); + vec_t v; + v.cast_load(v_smem + ((iter % num_smem_stages) * bdy + ty) * head_dim + tx * vec_size); + if (iter * bdy + ty < num_index_sets) { + float s = s_smem[(iter % bdx) * bdy + ty]; + st.merge(v, s, 1); + } + __syncwarp(); + cp_async::pred_load( + v_smem + ((iter % num_smem_stages) * bdy + ty) * head_dim + tx * vec_size, + V + partial_idx_to_offset((iter + num_smem_stages) * bdy + ty) * head_dim + + tx * vec_size, + (iter + num_smem_stages) * bdy + ty < num_index_sets); + cp_async::commit_group(); + } + cp_async::wait_group<0>(); + __syncwarp(); + + st.normalize(); + if constexpr (bdy > 1) { + warp_sync_state(st, v_smem, s_smem, tx, ty); + st.normalize(); + } + + st.o.cast_store(v_merged + merge_idx_to_offset() * head_dim + tx * vec_size); + if (s_merged != nullptr) { + s_merged[merge_idx_to_offset()] = st.get_lse(); + } + PROFILER_EVENT_END(profiler_closure, PersistentProfileEventType::kReduction); + } + } +}; + +template +cudaError_t BatchPagedAttentionPersistent(const Params params_1, const Params params_2, + const uint32_t num_blks_x, const uint32_t num_blks_y, + const cudaStream_t stream) { + using DTypeQ = typename Params::DTypeQ; + using DTypeKV = typename Params::DTypeKV; + using DTypeO = typename Params::DTypeO; + using IdType = typename Params::IdType; + constexpr uint32_t NUM_WARPS_Q_1 = get_num_warps_q(CTA_TILE_Q_1); + constexpr uint32_t NUM_WARPS_KV_1 = get_num_warps_kv(CTA_TILE_Q_1); + constexpr uint32_t NUM_MMA_Q_1 = get_num_mma_q(CTA_TILE_Q_1); + constexpr uint32_t NUM_MMA_KV_1 = 4; + constexpr uint32_t NUM_MMA_D_QK = HEAD_DIM_QK / 16; + constexpr uint32_t NUM_MMA_D_VO = HEAD_DIM_VO / 16; + using KTraits1 = KernelTraits; + constexpr uint32_t NUM_WARPS_Q_2 = get_num_warps_q(CTA_TILE_Q_2); + constexpr uint32_t NUM_WARPS_KV_2 = get_num_warps_kv(CTA_TILE_Q_2); + constexpr uint32_t NUM_MMA_Q_2 = get_num_mma_q(CTA_TILE_Q_2); + constexpr uint32_t NUM_MMA_KV_2 = 2; + using KTraits2 = KernelTraits; + + // Attention state reduction kernel + constexpr uint32_t NUM_THREADS = + KTraits1::NUM_THREADS > KTraits2::NUM_THREADS ? KTraits1::NUM_THREADS : KTraits2::NUM_THREADS; + using ReductionKTraits = + StateReductionKernelTraits; + size_t smem_size = + max(sizeof(typename KTraits1::SharedStorage), sizeof(typename KTraits2::SharedStorage)); + smem_size = max(smem_size, ReductionKTraits::SMEM_SIZE); + + // Launch persistent kernel + auto kernel = PersistentKernelTemplate, + BlockBatchPagedAttentionPersistent, + BlockBatchReductionPersistent>; + FLASHINFER_CUDA_CALL( + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + dim3 nblks(num_blks_x, num_blks_y); + dim3 nthrs(NUM_THREADS); + void* args[] = {(void*)¶ms_1, (void*)¶ms_2}; + FLASHINFER_CUDA_CALL( + cudaLaunchCooperativeKernel((void*)kernel, nblks, nthrs, args, smem_size, stream)); + return cudaSuccess; +} + +}; // namespace flashinfer + +#endif // FLASHINFER_PERSISTENT_CUH_ diff --git a/include/flashinfer/attention/persistent_template.cuh b/include/flashinfer/attention/persistent_template.cuh new file mode 100644 index 0000000000000000000000000000000000000000..3bd2331b3edc374d307243590aac86586f9b409d --- /dev/null +++ b/include/flashinfer/attention/persistent_template.cuh @@ -0,0 +1,100 @@ +/* + * Copyright (c) 2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef FLASHINFER_ATTENTION_PERSISTENT_TEMPLATE_CUH +#define FLASHINFER_ATTENTION_PERSISTENT_TEMPLATE_CUH + +#include +#include + +#include + +#include "../profiler.cuh" + +namespace flashinfer { +namespace cg = cooperative_groups; + +// Define profiler event types for persistent kernels +enum class PersistentProfileEventType { + kRunner1 = 0U, + kRunner2 = 1U, + kReduction = 2U, +}; + +struct ProfilerClosure { + PROFILER_CLOSURE_PARAMS_DECL +}; + +// Helper metafunction to find maximum threads among multiple BlockPersistentRunners +template +struct max_threads; + +template +struct max_threads { + static constexpr size_t value = Runner::KTraits::NUM_THREADS; +}; + +template +struct max_threads { + static constexpr size_t value = Runner1::KTraits::NUM_THREADS > Runner2::KTraits::NUM_THREADS + ? max_threads::value + : max_threads::value; +}; + +// Two runners version +template +__global__ __launch_bounds__( + max_threads:: + value) void PersistentKernelTemplate(const __grid_constant__ + typename BlockPersistentRunner1::Params params_1, + const __grid_constant__ + typename BlockPersistentRunner2::Params params_2) { + extern __shared__ uint8_t smem[]; + +#ifdef FLASHINFER_ENABLE_PROFILER + ProfilerClosure + profiler_closure; // no volatile as this is scope.CTA, and only threadIdx == 0 is modifying + PROFILER_INIT(params_1, smem, profiler_closure, 0, 1, (threadIdx.x == 0)); +#endif + + auto& smem_storage_1 = + reinterpret_cast(smem); + auto& smem_storage_2 = + reinterpret_cast(smem); + auto grid = cg::this_grid(); + +#ifndef FLASHINFER_ENABLE_PROFILER + BlockPersistentRunner1::Run(params_1, &smem_storage_1); + BlockPersistentRunner2::Run(params_2, &smem_storage_2); + + grid.sync(); + BlockReductionRunner::Run(params_1.partial_o, params_1.final_o, params_1.partial_lse, + params_1.final_lse, *(params_1.num_packed_qo_len), + params_1.gqa_group_size, params_1.num_kv_heads, params_1.merge_indptr, + params_1.merge_o_indices, smem); +#else + BlockPersistentRunner1::Run(params_1, &smem_storage_1, profiler_closure); + BlockPersistentRunner2::Run(params_2, &smem_storage_2, profiler_closure); + + grid.sync(); + BlockReductionRunner::Run(params_1.partial_o, params_1.final_o, params_1.partial_lse, + params_1.final_lse, *(params_1.num_packed_qo_len), + params_1.gqa_group_size, params_1.num_kv_heads, params_1.merge_indptr, + params_1.merge_o_indices, smem, profiler_closure); +#endif +} +} // namespace flashinfer + +#endif // FLASHINFER_ATTENTION_PERSISTENT_TEMPLATE_CUH diff --git a/include/flashinfer/attention/pod.cuh b/include/flashinfer/attention/pod.cuh new file mode 100644 index 0000000000000000000000000000000000000000..03ffdb755176c84b5744ec377f8f3a93a525f406 --- /dev/null +++ b/include/flashinfer/attention/pod.cuh @@ -0,0 +1,475 @@ +#ifndef FLASHINFER_POD_CUH_ +#define FLASHINFER_POD_CUH_ + +#include +#include +#include +#include +#include + +#include "../cp_async.cuh" +#include "../fastdiv.cuh" +#include "../frag_layout_swizzle.cuh" +#include "../layout.cuh" +#include "../math.cuh" +#include "../mma.cuh" +#include "../page.cuh" +#include "../permuted_smem.cuh" +#include "../pos_enc.cuh" +#include "../utils.cuh" +#include "cascade.cuh" +#include "decode.cuh" +#include "mask.cuh" +#include "prefill.cuh" +#include "variants.cuh" + +namespace flashinfer { + +namespace cg = cooperative_groups; +using cp_async::SharedMemFillMode; +using mma::MMAMode; + +enum Operation { + PREFILL = 0, + DECODE = 1, +}; + +template +__global__ __launch_bounds__(std::max( + KTraits_P::NUM_THREADS, + KTraits_D::NUM_THREADS)) void PODWithKVCacheTensorKernel(const uint32_t xsize, + const __grid_constant__ PrefillParams + prefill_params, + const __grid_constant__ DecodeParams + decode_params, + int* tbAssign) { + extern __shared__ uint8_t smem[]; + // PREFILL VARS + const uint32_t num_kv_heads_p = prefill_params.num_kv_heads; + const uint32_t num_chunks = prefill_params.partition_kv; + const uint32_t qo_len = prefill_params.qo_len; + + // DECODE VARS + const uint32_t padded_bsize = decode_params.padded_batch_size; + const uint32_t num_kv_heads_d = decode_params.paged_kv.num_heads; + + // THREADBLOCKS + const uint32_t prefill_blocks = num_kv_heads_p * xsize * (PartitionKV_P ? num_chunks : 1); + const uint32_t decode_blocks = padded_bsize * num_kv_heads_d; + + int op; + int linear_bid; + // SM-aware CTA scheduler + if (threadIdx.x == 0) { + // TODO_AK: If num_threads dont match, use virtual sub-CTAs. + // Requires changing block-level sync in main prefill/decode kernels. + constexpr int blk_factor_p = 1; + constexpr int blk_factor_d = 1; + + // SM-aware threadblock scheduler code + // Find out which SM this threadblock is scheduled on + int num_SMs; + // WARNING: nsmid has only been tested on A100/H100, and matches SM count + // No guarantee this will work on other GPUs + asm volatile("mov.u32 %0, %nsmid;" : "=r"(num_SMs)); + asm volatile("mov.u32 %0, %smid;" : "=r"(linear_bid)); + const int prefill_slots = (prefill_blocks + blk_factor_p - 1) / blk_factor_p; + const int decode_slots = (decode_blocks + blk_factor_d - 1) / blk_factor_d; + + if (prefill_slots <= decode_slots) { + // Total tags = (decode + prefill) / min(decode, prefill) + // = 1 + decode / prefill; when prefill < decode + const int total_tags = decode_slots / prefill_slots + 1; + // For this SM, what's the next operation we want to run? + op = (atomicAdd(&tbAssign[linear_bid], 1) % total_tags); + if (op > 0) { + op = 1; + } + } else { + // Total tags = (decode + prefill) / min(decode, prefill) + // = 1 + prefill / decode; when decode < prefill + const int pref_tags = prefill_slots / decode_slots; + + // For this SM, what's the next operation we want to run? + op = (atomicAdd(&tbAssign[linear_bid], 1) % (pref_tags + 1)); + if (op < pref_tags) { + op = 0; + } else { + op = 1; + } + } + + // Get the next blockId for that operation + linear_bid = atomicAdd(&tbAssign[num_SMs + op], 1); + // If the blockId obtained exceeds the max blockIds for that op, switch to the other op + if (op == 0 && linear_bid >= prefill_slots) { + linear_bid = atomicAdd(&tbAssign[num_SMs + 1], 1); + op = !op; + } else if (op == 1 && linear_bid >= decode_slots) { + op = !op; + linear_bid = atomicAdd(&tbAssign[num_SMs + 0], 1); + } + // Write the blockId and operation to shared memory + ((int*)smem)[0] = linear_bid; + ((int*)smem)[1] = op; + } + // Sync to wait for dynamic scheduler to finish + __syncthreads(); + // Fetch from shared memory the assigned blockId and operation. + linear_bid = ((int*)smem)[0]; + op = ((int*)smem)[1]; + // Sync to force all threads to wait + __syncthreads(); + + if (op == PREFILL) { + const uint32_t linear_tid = threadIdx.x; + // Return if threadId exceeds number of threads for this op + if (linear_tid >= 32 * KTraits_P::NUM_WARPS_Q * KTraits_P::NUM_WARPS_KV) return; + + const dim3 tid = dim3(linear_tid % 32, (linear_tid / 32) % KTraits_P::NUM_WARPS_Q, + (linear_tid / 32) / KTraits_P::NUM_WARPS_Q); + // dim3 nblks(ceil_div(qo_len * group_size, CTA_TILE_Q), 1, num_kv_heads); + // dim3 nblks(ceil_div(qo_len * group_size, CTA_TILE_Q), num_chunks, num_kv_heads); + // BlockID exceeds limit + if (linear_bid >= prefill_blocks) return; + + const uint32_t bx = linear_bid % xsize; + auto& smem_storage = reinterpret_cast(smem); + // Not partition_kv + if constexpr (!PartitionKV_P) { + const uint32_t chunk_idx = 0; + const uint32_t kv_head_idx = linear_bid / xsize; + SinglePrefillWithKVCacheDevice(prefill_params, smem_storage, tid, bx, chunk_idx, + kv_head_idx, 1, num_kv_heads_p); + } else { + const uint32_t chunk_idx = (linear_bid / xsize) % num_chunks; + const uint32_t kv_head_idx = linear_bid / (xsize * num_chunks); + SinglePrefillWithKVCacheDevice(prefill_params, smem_storage, tid, bx, chunk_idx, + kv_head_idx, num_chunks, num_kv_heads_p); + } + } else /* OP == DECODE */ { + auto& smem_storage = reinterpret_cast(smem); + // dim3 nblks_d(padded_batch_size_d, 1, num_kv_heads); + if (linear_bid >= decode_blocks) return; + + const uint32_t bx = linear_bid % padded_bsize; + const uint32_t kv_head_idx = linear_bid / padded_bsize; + + // dim3 nthrs_d(32, NUM_WARPS_Q_D, NUM_WARPS_KV_D); + const uint32_t linear_tid = threadIdx.x; + // Return if threadId exceeds number of threads for this op + if (linear_tid >= 32 * KTraits_D::NUM_WARPS_Q * KTraits_D::NUM_WARPS_KV) return; + + const dim3 tid = dim3(linear_tid % 32, (linear_tid / 32) % KTraits_D::NUM_WARPS_Q, + (linear_tid / 32) / KTraits_D::NUM_WARPS_Q); + + BatchPrefillWithPagedKVCacheDevice(decode_params, smem_storage, tid, bx, kv_head_idx, + num_kv_heads_d); + } +} + +template +cudaError_t PODWithKVCacheTensorDispatched(PrefillParams prefill_params, + typename PrefillParams::DTypeO* tmp_p, + DecodeParams decode_params, + typename DecodeParams::DTypeO* tmp_v, float* tmp_s, + bool enable_pdl, cudaStream_t stream) { + static_assert(std::is_same::value); + static_assert( + std::is_same::value); + static_assert(std::is_same::value); + // Ensure heads match + assert(prefill_params.num_kv_heads == decode_params.paged_kv.num_heads); + assert(prefill_params.num_qo_heads == decode_params.num_qo_heads); + // Prefill variable setup + using DTypeQ_P = typename PrefillParams::DTypeQ; + using DTypeKV_P = typename PrefillParams::DTypeKV; + using DTypeO_P = typename PrefillParams::DTypeO; + const uint32_t num_qo_heads = prefill_params.num_qo_heads; + const uint32_t num_kv_heads = prefill_params.num_kv_heads; + const uint32_t qo_len = prefill_params.qo_len; + const uint32_t kv_len = prefill_params.kv_len; + if (kv_len < qo_len && MASK_MODE_P == MaskMode::kCausal) { + std::ostringstream err_msg; + err_msg << "When mask_mode is set to MaskMode::kCausal, kv_len must be greater than or equal " + "to qo_len, got kv_len" + << kv_len << " and qo_len " << qo_len; + FLASHINFER_ERROR(err_msg.str()); + } + + const uint32_t group_size = num_qo_heads / num_kv_heads; + const uint_fastdiv group_size_fastdiv(group_size); + constexpr uint32_t NUM_MMA_D_QK = HEAD_DIM_QK / 16; + constexpr uint32_t NUM_MMA_D_VO = HEAD_DIM_VO / 16; + + uint32_t cta_tile_q_p = 0; + int64_t unpacked_qo_len = qo_len * group_size; + if (unpacked_qo_len > 64 && HEAD_DIM_VO < 256) { + cta_tile_q_p = 128; + } else { + auto compute_capacity = GetCudaComputeCapability(); + if (compute_capacity.first >= 8) { + // Ampere or newer + if (unpacked_qo_len > 16) { + // avg_packed_qo_len <= 64 + cta_tile_q_p = 64; + } else { + // avg_packed_qo_len <= 16 + cta_tile_q_p = 16; + } + } else { + // NOTE(Zihao): not enough shared memory on Turing for 1x4 warp layout + cta_tile_q_p = 64; + } + } + + // Decode vars setup + using DTypeQ_D = typename DecodeParams::DTypeQ; + using DTypeKV_D = typename DecodeParams::DTypeKV; + using DTypeO_D = typename DecodeParams::DTypeO; + const uint32_t padded_batch_size_d = decode_params.padded_batch_size; + constexpr uint32_t NUM_MMA_Q_D = get_num_mma_q(CTA_TILE_Q_D); + constexpr uint32_t NUM_WARPS_Q_D = get_num_warps_q(CTA_TILE_Q_D); + constexpr uint32_t NUM_WARPS_KV_D = get_num_warps_kv(CTA_TILE_Q_D); + + if (padded_batch_size_d == 0) { + // No request, skip + // this won't happen in CUDAGraph mode because we fixed the padded_batch_size + return cudaSuccess; + } + + // constexpr uint32_t NUM_MMA_D_QK = HEAD_DIM_QK / 16; + // constexpr uint32_t NUM_MMA_D_VO = HEAD_DIM_VO / 16; + using DTypeQKAccum_D = + typename std::conditional, half, + float>::type; + + int dev_id = 0; + FLASHINFER_CUDA_CALL(cudaGetDevice(&dev_id)); + int max_smem_per_sm = 0; + FLASHINFER_CUDA_CALL(cudaDeviceGetAttribute(&max_smem_per_sm, + cudaDevAttrMaxSharedMemoryPerMultiprocessor, dev_id)); + // we expect each sm execute two threadblocks + // TODO(Zihao): fix the following computation + const int num_ctas_per_sm = max_smem_per_sm > (16 * HEAD_DIM_QK * sizeof(DTypeQ_D) * 16) ? 2 : 1; + const int max_smem_per_threadblock = max_smem_per_sm / num_ctas_per_sm; + + constexpr uint32_t max_num_mma_kv_reg_d = + (HEAD_DIM_VO >= 128 && NUM_MMA_Q_D == 2 && POS_ENCODING_MODE == PosEncodingMode::kRoPELlama && + !USE_FP16_QK_REDUCTION) + ? 2 + : (8 / NUM_MMA_Q_D); + // TODO(Zihao): fix the following computation + const uint32_t max_num_mma_kv_smem_d = + (max_smem_per_threadblock / (16 * HEAD_DIM_QK * sizeof(DTypeQ_D)) - + NUM_MMA_Q_D * NUM_WARPS_Q_D) / + (2 * NUM_WARPS_KV_D); + + DISPATCH_CTA_TILE_Q(cta_tile_q_p, CTA_TILE_Q_P, { + constexpr uint32_t NUM_WARPS_Q_P = get_num_warps_q(CTA_TILE_Q_P); + constexpr uint32_t NUM_WARPS_KV_P = get_num_warps_kv(CTA_TILE_Q_P); + constexpr uint32_t NUM_MMA_Q_P = get_num_mma_q(CTA_TILE_Q_P); + + using DTypeQKAccum_P = + typename std::conditional, half, + float>::type; + + // we expect each sm execute two threadblocks + // TODO(Zihao): fix the following computation + const int num_ctas_per_sm_p = + max_smem_per_sm > (16 * HEAD_DIM_QK * sizeof(DTypeQ_P) * 16) ? 2 : 1; + const int max_smem_per_threadblock_p = max_smem_per_sm / num_ctas_per_sm_p; + + constexpr uint32_t max_num_mma_kv_reg_p = + (HEAD_DIM_VO >= 128 && NUM_MMA_Q_P == 2 && + POS_ENCODING_MODE == PosEncodingMode::kRoPELlama && !USE_FP16_QK_REDUCTION) + ? 2 + : (8 / NUM_MMA_Q_P); + // TODO(Zihao): fix the following computation + const uint32_t max_num_mma_kv_smem_p = + (max_smem_per_threadblock_p / (16 * HEAD_DIM_QK * sizeof(DTypeQ_P)) - + NUM_MMA_Q_P * NUM_WARPS_Q_P) / + (2 * NUM_WARPS_KV_P); + + // control NUM_MMA_KV for maximum warp occupancy + DISPATCH_NUM_MMA_KV(min(max_num_mma_kv_smem_p, max_num_mma_kv_reg_p), NUM_MMA_KV_P, { + using KTraits_P = + KernelTraits; + + if constexpr (KTraits_P::IsInvalid()) { + // Invalid configuration, skip + std::ostringstream err_msg; + err_msg << "FlashInfer Internal Error: Invalid configuration : NUM_MMA_Q=" << NUM_MMA_Q_P + << " NUM_MMA_D_QK=" << NUM_MMA_D_QK << " NUM_MMA_D_VO=" << NUM_MMA_D_VO + << " NUM_MMA_KV=" << NUM_MMA_KV_P << " NUM_WARPS_Q=" << NUM_WARPS_Q_P + << " NUM_WARPS_KV=" << NUM_WARPS_KV_P + << " please create an issue (https://github.com/flashinfer-ai/flashinfer/issues)" + " and report the issue to the developers."; + FLASHINFER_ERROR(err_msg.str()); + } else { + // Decode stuff + // TODO: Is there a way to avoid this nested dispatch? + DISPATCH_NUM_MMA_KV(min(max_num_mma_kv_smem_d, max_num_mma_kv_reg_d), NUM_MMA_KV_D, { + using KTraits_D = + KernelTraits; + if constexpr (KTraits_D::IsInvalid()) { + // Invalid configuration, skip + std::ostringstream err_msg; + err_msg + << "FlashInfer Internal Error: Invalid configuration : NUM_MMA_Q=" << NUM_MMA_Q_D + << " NUM_MMA_D_QK=" << NUM_MMA_D_QK << " NUM_MMA_D_VO=" << NUM_MMA_D_VO + << " NUM_MMA_KV=" << NUM_MMA_KV_D << " NUM_WARPS_Q=" << NUM_WARPS_Q_D + << " NUM_WARPS_KV=" << NUM_WARPS_KV_D + << " please create an issue (https://github.com/flashinfer-ai/flashinfer/issues)" + " and report the issue to the developers."; + FLASHINFER_ERROR(err_msg.str()); + } else { + // End decode stuff + constexpr uint32_t num_threads_p = (NUM_WARPS_Q_P * NUM_WARPS_KV_P) * WARP_SIZE; + size_t smem_size_p = sizeof(typename KTraits_P::SharedStorage); + size_t smem_size_d = sizeof(typename KTraits_D::SharedStorage); + + auto kernel = + PODWithKVCacheTensorKernel; + // Prefill: decide num_splits for split-kv + int num_blocks_per_sm = 0; + int num_sm = 0; + FLASHINFER_CUDA_CALL( + cudaDeviceGetAttribute(&num_sm, cudaDevAttrMultiProcessorCount, dev_id)); + // FLASHINFER_CUDA_CALL(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + // &num_blocks_per_sm, kernel, num_threads_p, smem_size_p)); + // Above function returns 0 for some reason, so we use a workaround + num_blocks_per_sm = std::max( + 1, std::min((int)(max_smem_per_sm / smem_size_p), (int)(256 / num_threads_p))); + uint32_t max_num_kv_chunks = + (num_blocks_per_sm * num_sm) / + (num_kv_heads * ceil_div(qo_len * group_size, KTraits_P::CTA_TILE_Q)); + uint32_t num_chunks; + if (max_num_kv_chunks > 0) { + uint32_t chunk_size = max(ceil_div(kv_len, max_num_kv_chunks), 256); + num_chunks = ceil_div(kv_len, chunk_size); + } else { + num_chunks = 0; + } + + // Setup new prefill params if (not) split + auto o_p = prefill_params.o; + auto lse_p = prefill_params.lse; + float* tmp_lse = (float*)(tmp_p + num_chunks * qo_len * num_qo_heads * HEAD_DIM_VO); + if (num_chunks <= 1 || tmp_p == nullptr) { + // Enough parallelism, do not split-kv + prefill_params.partition_kv = 0; + kernel = PODWithKVCacheTensorKernel; + } else { + // Use cooperative groups to increase occupancy + prefill_params.partition_kv = num_chunks; + prefill_params.o = tmp_p; + prefill_params.lse = tmp_lse; + kernel = PODWithKVCacheTensorKernel; + } + + // Setup new decode params if (not) split + auto o_d = decode_params.o; + auto lse_d = decode_params.lse; + if (tmp_v == nullptr) { + // do not partition kv + decode_params.partition_kv = false; + } else { + decode_params.partition_kv = true; + decode_params.o = tmp_v; + decode_params.lse = tmp_s; + } + uint32_t xsize = ceil_div(qo_len * group_size, KTraits_P::CTA_TILE_Q); + int nblks_p(xsize * (prefill_params.partition_kv ? prefill_params.partition_kv : 1) * + num_kv_heads); + int nthrs_p(32 * NUM_WARPS_Q_P * NUM_WARPS_KV_P); + + int nblks_d(padded_batch_size_d * 1 * num_kv_heads); + int nthrs_d(32 * NUM_WARPS_Q_D * NUM_WARPS_KV_D); + + // ******* Select final combined sizes here ******* / + size_t smem_size = max(smem_size_p, smem_size_d); + int nblks = nblks_p + nblks_d; + int nthrs = max(nthrs_p, nthrs_d); + + // printf("Smem: prefill %zu, decode %zu, total %zu\n", smem_size_p, smem_size_d, + // smem_size); printf("Blocks: prefill %d, decode %d, total %d\n", nblks_p, nblks_d, + // nblks); printf("Threads: prefill %d, decode %d, total %d\n", nthrs_p, nthrs_d, + // nthrs); + // ************************************************ / + + static int* tbAssign = nullptr; + if (tbAssign == nullptr) cudaMalloc(&tbAssign, sizeof(int) * (num_sm + 2)); + cudaMemset(tbAssign, 0, sizeof(int) * (num_sm + 2)); + + // Setup kernel arguments + void* args[] = {(void*)&xsize, (void*)&prefill_params, (void*)&decode_params, + (void*)&tbAssign}; + FLASHINFER_CUDA_CALL(cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + + // Launch kernel + if (enable_pdl) { + cudaLaunchAttribute attribute[1]; + cudaLaunchConfig_t config; + attribute[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attribute[0].val.programmaticStreamSerializationAllowed = 1; + config.attrs = attribute; + config.numAttrs = 1; + config.gridDim = nblks; + config.blockDim = nthrs; + config.dynamicSmemBytes = smem_size; + config.stream = stream; + FLASHINFER_CUDA_CALL(cudaLaunchKernelEx(&config, kernel, xsize, prefill_params, + decode_params, tbAssign)); + } else { + FLASHINFER_CUDA_CALL( + cudaLaunchKernel((void*)kernel, nblks, nthrs, args, smem_size, stream)); + } + + // Post-kernel stuff for split-kv prefill + if (!(num_chunks <= 1 || tmp_p == nullptr)) { + if constexpr (PrefillAttentionVariant::use_softmax) { + FLASHINFER_CUDA_CALL(MergeStates(tmp_p, tmp_lse, o_p, lse_p, num_chunks, qo_len, + num_qo_heads, HEAD_DIM_VO, stream)); + } else { + FLASHINFER_CUDA_CALL(AttentionSum(tmp_p, o_p, num_chunks, qo_len, num_qo_heads, + HEAD_DIM_VO, stream)); + } + } + // Post-kernel stuff for split-kv decode + if (tmp_v != nullptr) { + if constexpr (DecodeAttentionVariant::use_softmax) { + FLASHINFER_CUDA_CALL(VariableLengthMergeStates( + tmp_v, tmp_s, decode_params.merge_indptr, o_d, lse_d, + decode_params.max_total_num_rows, decode_params.total_num_rows, num_qo_heads, + HEAD_DIM_VO, enable_pdl, stream)); + } else { + FLASHINFER_CUDA_CALL(VariableLengthAttentionSum( + tmp_v, decode_params.merge_indptr, o_d, decode_params.max_total_num_rows, + decode_params.total_num_rows, num_qo_heads, HEAD_DIM_VO, enable_pdl, stream)); + } + } + } + }); + } + }); + }); + return cudaSuccess; +} + +} // namespace flashinfer + +#endif // FLASHINFER_PREFILL_CUH_ diff --git a/include/flashinfer/attention/prefill.cuh b/include/flashinfer/attention/prefill.cuh new file mode 100644 index 0000000000000000000000000000000000000000..b053b4268f6dd928174d7100704d297910000cde --- /dev/null +++ b/include/flashinfer/attention/prefill.cuh @@ -0,0 +1,2681 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef FLASHINFER_PREFILL_CUH_ +#define FLASHINFER_PREFILL_CUH_ + +#include +#include +#include +#include +#include + +#include "../cp_async.cuh" +#include "../fastdiv.cuh" +#ifdef FP16_QK_REDUCTION_SUPPORTED +#include "../fp16.h" +#endif +#include "../frag_layout_swizzle.cuh" +#include "../math.cuh" +#include "../mma.cuh" +#include "../page.cuh" +#include "../permuted_smem.cuh" +#include "../pos_enc.cuh" +#include "../utils.cuh" +#include "cascade.cuh" +#include "mask.cuh" +#include "variants.cuh" +namespace flashinfer { + +DEFINE_HAS_MEMBER(maybe_q_rope_offset) +DEFINE_HAS_MEMBER(maybe_k_rope_offset) +DEFINE_HAS_MEMBER(maybe_prefix_len_ptr) +DEFINE_HAS_MEMBER(maybe_token_pos_in_items_ptr) +DEFINE_HAS_MEMBER(token_pos_in_items_len) +DEFINE_HAS_MEMBER(maybe_max_item_len_ptr) + +namespace cg = cooperative_groups; +using cp_async::SharedMemFillMode; +using mma::MMAMode; + +constexpr uint32_t WARP_SIZE = 32; + +constexpr uint32_t get_num_warps_q(const uint32_t cta_tile_q) { + if (cta_tile_q > 16) { + return 4; + } else { + return 1; + } +} + +constexpr uint32_t get_num_warps_kv(const uint32_t cta_tile_kv) { + return 4 / get_num_warps_q(cta_tile_kv); +} + +constexpr uint32_t get_num_mma_q(const uint32_t cta_tile_q) { + if (cta_tile_q > 64) { + return 2; + } else { + return 1; + } +} + +template +struct SharedStorageQKVO { + union { + struct { + alignas(16) DTypeQ q_smem[CTA_TILE_Q * HEAD_DIM_QK]; + alignas(16) DTypeKV k_smem[CTA_TILE_KV * HEAD_DIM_QK]; + alignas(16) DTypeKV v_smem[CTA_TILE_KV * HEAD_DIM_VO]; + }; + struct { // NOTE(Zihao): synchronize attention states across warps + alignas( + 16) std::conditional_t cta_sync_o_smem; + alignas(16) std::conditional_t cta_sync_md_smem; + }; + alignas(16) DTypeO smem_o[CTA_TILE_Q * HEAD_DIM_VO]; + }; +}; + +template +struct KernelTraits { + static constexpr uint32_t NUM_STAGES = 1; // used for BatchAttention Template + static constexpr MaskMode MASK_MODE = MASK_MODE_; + static constexpr uint32_t NUM_MMA_Q = NUM_MMA_Q_; + static constexpr uint32_t NUM_MMA_KV = NUM_MMA_KV_; + static constexpr uint32_t NUM_MMA_D_QK = NUM_MMA_D_QK_; + static constexpr uint32_t NUM_MMA_D_VO = NUM_MMA_D_VO_; + static constexpr uint32_t NUM_WARPS_Q = NUM_WARPS_Q_; + static constexpr uint32_t NUM_WARPS_KV = NUM_WARPS_KV_; + static constexpr uint32_t NUM_THREADS = NUM_WARPS_Q * NUM_WARPS_KV * WARP_SIZE; + static constexpr uint32_t NUM_WARPS = NUM_WARPS_Q * NUM_WARPS_KV; + static constexpr uint32_t HEAD_DIM_QK = NUM_MMA_D_QK * 16; + static constexpr uint32_t HEAD_DIM_VO = NUM_MMA_D_VO * 16; + static constexpr uint32_t UPCAST_STRIDE_Q = HEAD_DIM_QK / upcast_size(); + static constexpr uint32_t UPCAST_STRIDE_K = HEAD_DIM_QK / upcast_size(); + static constexpr uint32_t UPCAST_STRIDE_V = HEAD_DIM_VO / upcast_size(); + static constexpr uint32_t UPCAST_STRIDE_O = HEAD_DIM_VO / upcast_size(); + static constexpr uint32_t CTA_TILE_Q = CTA_TILE_Q_; + static constexpr uint32_t CTA_TILE_KV = NUM_MMA_KV * NUM_WARPS_KV * 16; + + static constexpr SwizzleMode SWIZZLE_MODE_Q = SwizzleMode::k128B; + static constexpr SwizzleMode SWIZZLE_MODE_KV = + (sizeof(DTypeKV_) == 1 && HEAD_DIM_VO == 64) ? SwizzleMode::k64B : SwizzleMode::k128B; + static constexpr uint32_t KV_THR_LAYOUT_ROW = SWIZZLE_MODE_KV == SwizzleMode::k128B ? 4 : 8; + static constexpr uint32_t KV_THR_LAYOUT_COL = SWIZZLE_MODE_KV == SwizzleMode::k128B ? 8 : 4; + static constexpr PosEncodingMode POS_ENCODING_MODE = POS_ENCODING_MODE_; + using DTypeQ = DTypeQ_; + using DTypeKV = DTypeKV_; + using DTypeO = DTypeO_; + using DTypeQKAccum = DTypeQKAccum_; + using IdType = IdType_; + using AttentionVariant = AttentionVariant_; + + static constexpr bool IsInvalid() { + return ((NUM_MMA_D_VO < 4) || (NUM_MMA_D_VO == 4 && NUM_MMA_KV % 2 == 1) || + (POS_ENCODING_MODE == PosEncodingMode::kRoPELlama && NUM_MMA_D_VO > 4 && + NUM_MMA_D_VO % (2 * NUM_WARPS_Q) != 0) || + (NUM_MMA_Q * (8 * NUM_MMA_D_VO + 2 * sizeof(DTypeQKAccum) * NUM_MMA_KV) >= 256) || + (sizeof(DTypeKV) == 1 && NUM_MMA_KV * 2 % NUM_WARPS_Q != 0) || + (sizeof(DTypeKV) == 1 && POS_ENCODING_MODE == PosEncodingMode::kRoPELlama)); + } + + using SharedStorage = SharedStorageQKVO; +#ifdef FP16_QK_REDUCTION_SUPPORTED + template + static constexpr DT getNegInf() { + if constexpr (std::is_same::value) { + return std::bit_cast(fp16_ieee_from_fp32_value(-math::inf)); + } else { + return static_cast(-math::inf); + } + } + + static constexpr DTypeQKAccum MaskFillValue = + AttentionVariant::use_softmax ? getNegInf() : DTypeQKAccum(0.f); +#else + static_assert(!std::is_same::value, + "Set -DFP16_QK_REDUCTION_SUPPORTED and install boost_math " + "then recompile to support fp16 reduction"); + static constexpr DTypeQKAccum MaskFillValue = + AttentionVariant::use_softmax ? DTypeQKAccum(-math::inf) : DTypeQKAccum(0.f); +#endif +}; + +namespace { + +template +__device__ __forceinline__ uint32_t get_warp_idx_q(const uint32_t tid_y = threadIdx.y) { + if constexpr (KTraits::NUM_WARPS_Q == 1) { + return 0; + } else { + return tid_y; + } +} + +template +__device__ __forceinline__ uint32_t get_warp_idx_kv(const uint32_t tid_z = threadIdx.z) { + if constexpr (KTraits::NUM_WARPS_KV == 1) { + return 0; + } else { + return tid_z; + } +} + +template +__device__ __forceinline__ uint32_t get_warp_idx(const uint32_t tid_y = threadIdx.y, + const uint32_t tid_z = threadIdx.z) { + return get_warp_idx_kv(tid_z) * KTraits::NUM_WARPS_Q + get_warp_idx_q(tid_y); +} + +/*! + * \brief Apply Llama style rotary embedding to two 16x16 fragments. + * \tparam T The data type of the input fragments. + * \param x_first_half First fragment x[offset:offset+16, j*16:(j+1)*16] + * \param x_second_half Second fragment x[offset:offset*16, j*16+d/2:(j+1)*16+d/2] + * \param rope_freq Rope frequency + * \param offset The offset of the first row in both fragments. + * \note The sin/cos computation is slow, especially for A100 GPUs which has low + * non tensor-ops flops, will optimize in the future. + */ +template +__device__ __forceinline__ void k_frag_apply_llama_rope(T* x_first_half, T* x_second_half, + const float* rope_freq, + const uint32_t kv_offset) { + static_assert(sizeof(T) == 2); +#pragma unroll + for (uint32_t reg_id = 0; reg_id < 8; ++reg_id) { + float cos, sin, tmp; + // 0 1 | 2 3 + // --------- + // 4 5 | 6 7 + uint32_t i = reg_id / 4, j = (reg_id % 4) / 2; + __sincosf(float(kv_offset + 8 * i) * rope_freq[2 * j + reg_id % 2], &sin, &cos); + tmp = x_first_half[reg_id]; + x_first_half[reg_id] = (tmp * cos - (float)x_second_half[reg_id] * sin); + x_second_half[reg_id] = ((float)x_second_half[reg_id] * cos + tmp * sin); + } +} + +template +__device__ __forceinline__ void q_frag_apply_llama_rope(T* x_first_half, T* x_second_half, + const float* rope_freq, + const uint32_t qo_packed_offset, + const uint_fastdiv group_size) { +#pragma unroll + for (uint32_t reg_id = 0; reg_id < 8; ++reg_id) { + float cos, sin, tmp; + // 0 1 | 4 5 + // --------- + // 2 3 | 6 7 + uint32_t i = ((reg_id % 4) / 2), j = (reg_id / 4); + __sincosf(float((qo_packed_offset + 8 * i) / group_size) * rope_freq[2 * j + reg_id % 2], &sin, + &cos); + tmp = x_first_half[reg_id]; + x_first_half[reg_id] = (tmp * cos - (float)x_second_half[reg_id] * sin); + x_second_half[reg_id] = ((float)x_second_half[reg_id] * cos + tmp * sin); + } +} + +template +__device__ __forceinline__ void q_frag_apply_llama_rope_with_pos(T* x_first_half, T* x_second_half, + const float* rope_freq, + const uint32_t qo_packed_offset, + const uint_fastdiv group_size, + const IdType* q_rope_offset) { + float pos[2] = {static_cast(q_rope_offset[qo_packed_offset / group_size]), + static_cast(q_rope_offset[(qo_packed_offset + 8) / group_size])}; +#pragma unroll + for (uint32_t reg_id = 0; reg_id < 8; ++reg_id) { + float cos, sin, tmp; + // 0 1 | 4 5 + // --------- + // 2 3 | 6 7 + uint32_t i = ((reg_id % 4) / 2), j = (reg_id / 4); + __sincosf(pos[i] * rope_freq[2 * j + reg_id % 2], &sin, &cos); + tmp = x_first_half[reg_id]; + x_first_half[reg_id] = (tmp * cos - (float)x_second_half[reg_id] * sin); + x_second_half[reg_id] = ((float)x_second_half[reg_id] * cos + tmp * sin); + } +} + +/*! + * \brief Produce k/v fragments from global memory to shared memory. + * \tparam fill_mode The fill mode of the shared memory. + * \tparam NUM_MMA_D_VO The number of fragments in y dimension. + * \tparam NUM_MMA_KV The number of fragments in z dimension. + * \tparam num_warps The number of warps in the threadblock. + * \tparam T The data type of the input tensor. + * \param smem The shared memory to store kv fragments. + * \param gptr The global memory pointer. + * \param kv_idx_base The base kv index. + * \param kv_len The length of kv tensor. + */ +template +__device__ __forceinline__ void produce_kv(smem_t smem, + uint32_t* smem_offset, typename KTraits::DTypeKV** gptr, + const uint32_t stride_n, const uint32_t kv_idx_base, + const uint32_t kv_len, const dim3 tid = threadIdx) { + // NOTE: for fp8, this function doesn't work for head_dim = 64 at the moment + using DTypeKV = typename KTraits::DTypeKV; + constexpr uint32_t CTA_TILE_KV = KTraits::CTA_TILE_KV; + constexpr uint32_t NUM_WARPS = KTraits::NUM_WARPS; + constexpr uint32_t NUM_WARPS_Q = KTraits::NUM_WARPS_Q; + constexpr uint32_t NUM_MMA_D = produce_v ? KTraits::NUM_MMA_D_VO : KTraits::NUM_MMA_D_QK; + constexpr uint32_t NUM_MMA_KV = KTraits::NUM_MMA_KV; + constexpr uint32_t UPCAST_STRIDE = + produce_v ? KTraits::UPCAST_STRIDE_V : KTraits::UPCAST_STRIDE_K; + const uint32_t warp_idx = get_warp_idx(tid.y, tid.z), lane_idx = tid.x; + + if constexpr (KTraits::SWIZZLE_MODE_KV == SwizzleMode::k128B) { + uint32_t kv_idx = kv_idx_base + warp_idx * 4 + lane_idx / 8; + // NOTE: NUM_MMA_KV * 4 / NUM_WARPS_Q = NUM_WARPS_KV * NUM_MMA_KV * 4 / num_warps + static_assert(NUM_MMA_KV * 4 % NUM_WARPS_Q == 0); +#pragma unroll + for (uint32_t i = 0; i < NUM_MMA_KV * 4 / NUM_WARPS_Q; ++i) { +#pragma unroll + for (uint32_t j = 0; j < NUM_MMA_D / (8 / sizeof(DTypeKV)); ++j) { + smem.load_128b_async(*smem_offset, *gptr, kv_idx < kv_len); + *smem_offset = smem.template advance_offset_by_column<8>(*smem_offset, j); + *gptr += 8 * upcast_size(); + } + kv_idx += NUM_WARPS * 4; + *smem_offset = + smem.template advance_offset_by_row(*smem_offset) - + sizeof(DTypeKV) * NUM_MMA_D; + *gptr += NUM_WARPS * 4 * stride_n - sizeof(DTypeKV) * NUM_MMA_D * upcast_size(); + } + *smem_offset -= CTA_TILE_KV * UPCAST_STRIDE; + } else { + uint32_t kv_idx = kv_idx_base + warp_idx * 8 + lane_idx / 4; + // NOTE: NUM_MMA_KV * 2 / NUM_WARPS_Q = NUM_WARPS_KV * NUM_MMA_KV * 2 / num_warps + static_assert(NUM_MMA_KV * 2 % NUM_WARPS_Q == 0); +#pragma unroll + for (uint32_t i = 0; i < NUM_MMA_KV * 2 / NUM_WARPS_Q; ++i) { + smem.load_128b_async(*smem_offset, *gptr, kv_idx < kv_len); + *smem_offset = + smem.template advance_offset_by_row(*smem_offset); + kv_idx += NUM_WARPS * 8; + *gptr += NUM_WARPS * 8 * stride_n; + } + *smem_offset -= KTraits::CTA_TILE_KV * UPCAST_STRIDE; + } +} + +template +__device__ __forceinline__ void page_produce_kv(typename KTraits::SharedStorage* smem_storage, + uint32_t* smem_offset, + typename KTraits::DTypeKV* kv_ptr, + const uint32_t kv_idx_base, + const size_t* thr_local_kv_offset, + const uint32_t kv_len, const uint32_t warp_idx, + const uint32_t lane_idx) { + // NOTE: for fp8, this function doesn't work for head_dim = 64 at the moment + smem_t smem(produce_v ? smem_storage->v_smem : smem_storage->k_smem); + using DType = typename KTraits::DTypeKV; + using IdType = typename KTraits::IdType; + constexpr SharedMemFillMode fill_mode = + produce_v ? SharedMemFillMode::kFillZero : SharedMemFillMode::kNoFill; + constexpr uint32_t NUM_WARPS = KTraits::NUM_WARPS; + constexpr uint32_t NUM_WARPS_Q = KTraits::NUM_WARPS_Q; + constexpr uint32_t NUM_MMA_KV = KTraits::NUM_MMA_KV; + constexpr uint32_t NUM_MMA_D = produce_v ? KTraits::NUM_MMA_D_VO : KTraits::NUM_MMA_D_QK; + constexpr uint32_t UPCAST_STRIDE = + produce_v ? KTraits::UPCAST_STRIDE_V : KTraits::UPCAST_STRIDE_K; + if constexpr (KTraits::SWIZZLE_MODE_KV == SwizzleMode::k128B) { + uint32_t kv_idx = kv_idx_base + warp_idx * 4 + lane_idx / 8; + // NOTE: NUM_MMA_KV * 4 / NUM_WARPS_Q = NUM_WARPS_KV * NUM_MMA_KV * 4 / num_warps + static_assert(NUM_MMA_KV * 4 % NUM_WARPS_Q == 0); +#pragma unroll + for (uint32_t i = 0; i < NUM_MMA_KV * 4 / NUM_WARPS_Q; ++i) { + DType* gptr = kv_ptr + thr_local_kv_offset[i]; +#pragma unroll + for (uint32_t j = 0; j < NUM_MMA_D / (8 / sizeof(DType)); ++j) { + smem.load_128b_async(*smem_offset, gptr, kv_idx < kv_len); + *smem_offset = smem.template advance_offset_by_column<8>(*smem_offset, j); + gptr += 8 * upcast_size(); + } + kv_idx += NUM_WARPS * 4; + *smem_offset = + smem.template advance_offset_by_row(*smem_offset) - + sizeof(DType) * NUM_MMA_D; + } + *smem_offset -= KTraits::CTA_TILE_KV * UPCAST_STRIDE; + } else { + uint32_t kv_idx = kv_idx_base + warp_idx * 8 + lane_idx / 4; + // NOTE: NUM_MMA_KV * 2 / NUM_WARPS_Q = NUM_WARPS_KV * NUM_MMA_KV * 2 / num_warps + static_assert(NUM_MMA_KV * 2 % NUM_WARPS_Q == 0); +#pragma unroll + for (uint32_t i = 0; i < NUM_MMA_KV * 2 / NUM_WARPS_Q; ++i) { + DType* gptr = kv_ptr + thr_local_kv_offset[i]; + smem.load_128b_async(*smem_offset, gptr, kv_idx < kv_len); + kv_idx += NUM_WARPS * 8; + *smem_offset = + smem.template advance_offset_by_row(*smem_offset); + } + *smem_offset -= KTraits::CTA_TILE_KV * UPCAST_STRIDE; + } +} + +template +__device__ __forceinline__ void init_rope_freq(float (*rope_freq)[4], const float rope_rcp_scale, + const float rope_rcp_theta, + const uint32_t tid_x = threadIdx.x) { + constexpr uint32_t HEAD_DIM = KTraits::NUM_MMA_D_QK * 16; + const uint32_t lane_idx = tid_x; +#pragma unroll + for (uint32_t mma_d = 0; mma_d < KTraits::NUM_MMA_D_VO / 2; ++mma_d) { +#pragma unroll + for (uint32_t j = 0; j < 4; ++j) { + rope_freq[mma_d][j] = + rope_rcp_scale * + __powf(rope_rcp_theta, + float(2 * ((mma_d * 16 + (j / 2) * 8 + (lane_idx % 4) * 2 + (j % 2)) % + (HEAD_DIM / 2))) / + float(HEAD_DIM)); + } + } +} + +template +__device__ __forceinline__ void init_states(typename KTraits::AttentionVariant variant, + float (*o_frag)[KTraits::NUM_MMA_D_VO][8], + typename KTraits::DTypeQKAccum (*m)[2], float (*d)[2]) { +#pragma unroll + for (uint32_t mma_q = 0; mma_q < KTraits::NUM_MMA_Q; ++mma_q) { +#pragma unroll + for (uint32_t mma_d = 0; mma_d < KTraits::NUM_MMA_D_VO; ++mma_d) { +#pragma unroll + for (uint32_t reg_id = 0; reg_id < 8; ++reg_id) { + o_frag[mma_q][mma_d][reg_id] = 0.f; + } + } + } + + if constexpr (variant.use_softmax) { +#pragma unroll + for (uint32_t mma_q = 0; mma_q < KTraits::NUM_MMA_Q; ++mma_q) { +#pragma unroll + for (uint32_t j = 0; j < 2; ++j) { + m[mma_q][j] = typename KTraits::DTypeQKAccum(-math::inf); + d[mma_q][j] = 1.f; + } + } + } +} + +template +__device__ __forceinline__ void load_q_global_smem( + uint32_t packed_offset, const uint32_t qo_upper_bound, typename KTraits::DTypeQ* q_ptr_base, + const uint32_t q_stride_n, const uint32_t q_stride_h, const uint_fastdiv group_size, + smem_t* q_smem, const dim3 tid = threadIdx) { + using DTypeQ = typename KTraits::DTypeQ; + constexpr uint32_t UPCAST_STRIDE_Q = KTraits::UPCAST_STRIDE_Q; + const uint32_t lane_idx = tid.x, warp_idx_x = get_warp_idx_q(tid.y); + + if (get_warp_idx_kv(tid.z) == 0) { + uint32_t q_smem_offset_w = q_smem->get_permuted_offset( + warp_idx_x * KTraits::NUM_MMA_Q * 16 + lane_idx / 8, lane_idx % 8); + +#pragma unroll + for (uint32_t mma_q = 0; mma_q < KTraits::NUM_MMA_Q; ++mma_q) { +#pragma unroll + for (uint32_t j = 0; j < 2 * 2; ++j) { + uint32_t q, r; + group_size.divmod(packed_offset + lane_idx / 8 + mma_q * 16 + j * 4, q, r); + const uint32_t q_idx = q; + DTypeQ* q_ptr = + q_ptr_base + q * q_stride_n + r * q_stride_h + (lane_idx % 8) * upcast_size(); +#pragma unroll + for (uint32_t mma_do = 0; mma_do < KTraits::NUM_MMA_D_QK / 4; ++mma_do) { + // load q fragment from gmem to smem + q_smem->load_128b_async(q_smem_offset_w, q_ptr, + q_idx < qo_upper_bound); + q_smem_offset_w = q_smem->template advance_offset_by_column<8>(q_smem_offset_w, mma_do); + q_ptr += 8 * upcast_size(); + } + q_smem_offset_w = + q_smem->template advance_offset_by_row<4, UPCAST_STRIDE_Q>(q_smem_offset_w) - + 2 * KTraits::NUM_MMA_D_QK; + } + } + } +} + +template +__device__ __forceinline__ void q_smem_inplace_apply_rotary( + const uint32_t q_packed_idx, const uint32_t qo_len, const uint32_t kv_len, + const uint_fastdiv group_size, smem_t* q_smem, + uint32_t* q_smem_offset_r, float (*rope_freq)[4], const dim3 tid = threadIdx) { + if (get_warp_idx_kv(tid.z) == 0) { + constexpr uint32_t UPCAST_STRIDE_Q = KTraits::UPCAST_STRIDE_Q; + const uint32_t lane_idx = tid.x; + uint32_t q_frag_local[2][4]; + static_assert(KTraits::NUM_MMA_D_QK % 4 == 0, "NUM_MMA_D_QK must be a multiple of 4"); +#pragma unroll + for (uint32_t mma_q = 0; mma_q < KTraits::NUM_MMA_Q; ++mma_q) { + uint32_t q_smem_offset_r_first_half = *q_smem_offset_r; +#pragma unroll + for (uint32_t mma_di = 0; mma_di < KTraits::NUM_MMA_D_QK / 2; ++mma_di) { + q_smem->ldmatrix_m8n8x4(q_smem_offset_r_first_half, q_frag_local[0]); + uint32_t q_smem_offset_r_last_half = + q_smem->template advance_offset_by_column( + q_smem_offset_r_first_half, 0); + q_smem->ldmatrix_m8n8x4(q_smem_offset_r_last_half, q_frag_local[1]); + q_frag_apply_llama_rope( + (typename KTraits::DTypeQ*)q_frag_local[0], (typename KTraits::DTypeQ*)q_frag_local[1], + rope_freq[mma_di], + q_packed_idx + kv_len * group_size - qo_len * group_size + mma_q * 16 + lane_idx / 4, + group_size); + q_smem->stmatrix_m8n8x4(q_smem_offset_r_last_half, q_frag_local[1]); + q_smem->stmatrix_m8n8x4(q_smem_offset_r_first_half, q_frag_local[0]); + q_smem_offset_r_first_half = + q_smem->template advance_offset_by_column<2>(q_smem_offset_r_first_half, mma_di); + } + *q_smem_offset_r += 16 * UPCAST_STRIDE_Q; + } + *q_smem_offset_r -= KTraits::NUM_MMA_Q * 16 * UPCAST_STRIDE_Q; + } +} + +template +__device__ __forceinline__ void q_smem_inplace_apply_rotary_with_pos( + const uint32_t q_packed_idx_base, const typename KTraits::IdType* q_rope_offset, + smem_t* q_smem, const uint_fastdiv group_size, + uint32_t* q_smem_offset_r, float (*rope_freq)[4], const dim3 tid = threadIdx) { + if (get_warp_idx_kv(tid.z) == 0) { + constexpr uint32_t UPCAST_STRIDE_Q = KTraits::UPCAST_STRIDE_Q; + const uint32_t lane_idx = tid.x; + uint32_t q_frag_local[2][4]; + static_assert(KTraits::NUM_MMA_D_QK % 4 == 0, "NUM_MMA_D_QK must be a multiple of 4"); +#pragma unroll + for (uint32_t mma_q = 0; mma_q < KTraits::NUM_MMA_Q; ++mma_q) { + uint32_t q_smem_offset_r_first_half = *q_smem_offset_r; +#pragma unroll + for (uint32_t mma_di = 0; mma_di < KTraits::NUM_MMA_D_QK / 2; ++mma_di) { + q_smem->ldmatrix_m8n8x4(q_smem_offset_r_first_half, q_frag_local[0]); + uint32_t q_smem_offset_r_last_half = + q_smem->template advance_offset_by_column( + q_smem_offset_r_first_half, 0); + q_smem->ldmatrix_m8n8x4(q_smem_offset_r_last_half, q_frag_local[1]); + q_frag_apply_llama_rope_with_pos( + (typename KTraits::DTypeQ*)q_frag_local[0], (typename KTraits::DTypeQ*)q_frag_local[1], + rope_freq[mma_di], q_packed_idx_base + mma_q * 16 + lane_idx / 4, group_size, + q_rope_offset); + q_smem->stmatrix_m8n8x4(q_smem_offset_r_last_half, q_frag_local[1]); + q_smem->stmatrix_m8n8x4(q_smem_offset_r_first_half, q_frag_local[0]); + q_smem_offset_r_first_half = + q_smem->template advance_offset_by_column<2>(q_smem_offset_r_first_half, mma_di); + } + *q_smem_offset_r += 16 * UPCAST_STRIDE_Q; + } + *q_smem_offset_r -= KTraits::NUM_MMA_Q * 16 * UPCAST_STRIDE_Q; + } +} + +template +__device__ __forceinline__ void k_smem_inplace_apply_rotary( + const uint32_t kv_idx_base, smem_t* k_smem, uint32_t* k_smem_offset_r, + float (*rope_freq)[4], const dim3 tid = threadIdx) { + using DTypeKV = typename KTraits::DTypeKV; + static_assert(sizeof(DTypeKV) == 2); + constexpr uint32_t UPCAST_STRIDE_K = KTraits::UPCAST_STRIDE_K; + uint32_t k_frag_local[2][4]; + const uint32_t lane_idx = tid.x; + if constexpr (KTraits::NUM_MMA_D_QK == 4 && KTraits::NUM_WARPS_Q == 4) { + static_assert(KTraits::NUM_WARPS_KV == 1); + const uint32_t warp_idx = get_warp_idx_q(tid.y); + // horizontal-axis: y + // vertical-axis: z + // | 1-16 | 16-32 | 32-48 | 48-64 | + // | 1-16 | warp_idx=0 | warp_idx=1 | warp_idx=0 | warp_idx=1 | + // | 16-32 | warp_idx=2 | warp_idx=3 | warp_idx=2 | warp_idx=3 | + static_assert(KTraits::NUM_MMA_KV % 2 == 0, + "when NUM_MMA_D_QK == 4, NUM_MMA_KV must be a multiple of 2"); + uint32_t kv_idx = kv_idx_base + (warp_idx / 2) * 16 + lane_idx / 4; + *k_smem_offset_r = + (*k_smem_offset_r ^ (0x2 * (warp_idx % 2))) + (warp_idx / 2) * 16 * UPCAST_STRIDE_K; +#pragma unroll + for (uint32_t i = 0; i < KTraits::NUM_MMA_KV / 2; ++i) { + uint32_t k_smem_offset_r_first_half = *k_smem_offset_r; + uint32_t mma_di = (warp_idx % 2); + k_smem->ldmatrix_m8n8x4(k_smem_offset_r_first_half, k_frag_local[0]); + uint32_t k_smem_offset_r_last_half = + k_smem->template advance_offset_by_column<4>(k_smem_offset_r_first_half, 0); + k_smem->ldmatrix_m8n8x4(k_smem_offset_r_last_half, k_frag_local[1]); + k_frag_apply_llama_rope((DTypeKV*)k_frag_local[0], (DTypeKV*)k_frag_local[1], + rope_freq[mma_di], kv_idx); + k_smem->stmatrix_m8n8x4(k_smem_offset_r_last_half, k_frag_local[1]); + k_smem->stmatrix_m8n8x4(k_smem_offset_r_first_half, k_frag_local[0]); + *k_smem_offset_r += 32 * UPCAST_STRIDE_K; + kv_idx += 32; + } + *k_smem_offset_r = (*k_smem_offset_r ^ (0x2 * (warp_idx % 2))) - + ((warp_idx / 2) + KTraits::NUM_MMA_KV) * 16 * UPCAST_STRIDE_K; + } else { + const uint32_t warp_idx_x = get_warp_idx_q(tid.y), + warp_idx_z = get_warp_idx_kv(tid.z); + static_assert(KTraits::NUM_MMA_D_QK % (2 * KTraits::NUM_WARPS_Q) == 0); + // horizontal axis: y + // vertical axis: z + // | (warp_idx_z, warp_idx_x) | 1-16 | 16-32 | 32-48 | 48-64 | ... + // | 1-16*NUM_MMA_KV | (0, 0) | (0, 1) | (0, 2) | (0, 3) | ... + // | 16*NUM_MMA_KV-32*NUM_MMA_KV | (1, 0) | (1, 1) | (1, 2) | (1, 3) | ... + // ... + uint32_t kv_idx = kv_idx_base + (warp_idx_z * KTraits::NUM_MMA_KV * 16) + lane_idx / 4; + *k_smem_offset_r = *k_smem_offset_r ^ (0x2 * warp_idx_x); +#pragma unroll + for (uint32_t i = 0; i < KTraits::NUM_MMA_KV; ++i) { + uint32_t k_smem_offset_r_first_half = *k_smem_offset_r; +#pragma unroll + for (uint32_t j = 0; j < KTraits::NUM_MMA_D_QK / (2 * KTraits::NUM_WARPS_Q); ++j) { + uint32_t mma_di = warp_idx_x + j * KTraits::NUM_WARPS_Q; + k_smem->ldmatrix_m8n8x4(k_smem_offset_r_first_half, k_frag_local[0]); + uint32_t k_smem_offset_r_last_half = + k_smem->template advance_offset_by_column( + k_smem_offset_r_first_half, 0); + k_smem->ldmatrix_m8n8x4(k_smem_offset_r_last_half, k_frag_local[1]); + k_frag_apply_llama_rope((DTypeKV*)k_frag_local[0], (DTypeKV*)k_frag_local[1], + rope_freq[mma_di], kv_idx); + k_smem->stmatrix_m8n8x4(k_smem_offset_r_last_half, k_frag_local[1]); + k_smem->stmatrix_m8n8x4(k_smem_offset_r_first_half, k_frag_local[0]); + k_smem_offset_r_first_half = + k_smem->template advance_offset_by_column<2 * KTraits::NUM_WARPS_Q>( + k_smem_offset_r_first_half, mma_di); + } + *k_smem_offset_r += 16 * UPCAST_STRIDE_K; + kv_idx += 16; + } + *k_smem_offset_r = + (*k_smem_offset_r ^ (0x2 * warp_idx_x)) - KTraits::NUM_MMA_KV * 16 * UPCAST_STRIDE_K; + } +} + +template +__device__ __forceinline__ void compute_qk( + smem_t* q_smem, uint32_t* q_smem_offset_r, + smem_t* k_smem, uint32_t* k_smem_offset_r, + typename KTraits::DTypeQKAccum (*s_frag)[KTraits::NUM_MMA_KV][8]) { + constexpr uint32_t UPCAST_STRIDE_Q = KTraits::UPCAST_STRIDE_Q; + constexpr uint32_t UPCAST_STRIDE_K = KTraits::UPCAST_STRIDE_K; + uint32_t a_frag[KTraits::NUM_MMA_Q][4], b_frag[4]; + // compute q*k^T +#pragma unroll + for (uint32_t mma_d = 0; mma_d < KTraits::NUM_MMA_D_QK; ++mma_d) { +#pragma unroll + for (uint32_t mma_q = 0; mma_q < KTraits::NUM_MMA_Q; ++mma_q) { + q_smem->ldmatrix_m8n8x4(*q_smem_offset_r, a_frag[mma_q]); + *q_smem_offset_r = + q_smem->template advance_offset_by_row<16, UPCAST_STRIDE_Q>(*q_smem_offset_r); + } + + *q_smem_offset_r = q_smem->template advance_offset_by_column<2>(*q_smem_offset_r, mma_d) - + KTraits::NUM_MMA_Q * 16 * UPCAST_STRIDE_Q; + +#pragma unroll + for (uint32_t mma_kv = 0; mma_kv < KTraits::NUM_MMA_KV; ++mma_kv) { + if constexpr (sizeof(typename KTraits::DTypeKV) == 1) { + uint32_t b_frag_f8[2]; + if (mma_d % 2 == 0) { + k_smem->ldmatrix_m8n8x4_left_half(*k_smem_offset_r, b_frag_f8); + } else { + k_smem->ldmatrix_m8n8x4_right_half(*k_smem_offset_r, b_frag_f8); + } + b_frag_f8[0] = frag_layout_swizzle_16b_to_8b(b_frag_f8[0]); + b_frag_f8[1] = frag_layout_swizzle_16b_to_8b(b_frag_f8[1]); + vec_cast::cast<8>( + (typename KTraits::DTypeQ*)b_frag, (typename KTraits::DTypeKV*)b_frag_f8); + } else { + k_smem->ldmatrix_m8n8x4(*k_smem_offset_r, b_frag); + } + *k_smem_offset_r = + k_smem->template advance_offset_by_row<16, UPCAST_STRIDE_K>(*k_smem_offset_r); + +#pragma unroll + for (uint32_t mma_q = 0; mma_q < KTraits::NUM_MMA_Q; ++mma_q) { + if constexpr (std::is_same_v) { + if (mma_d == 0) { + mma::mma_sync_m16n16k16_row_col_f16f16f32( + s_frag[mma_q][mma_kv], a_frag[mma_q], b_frag); + } else { + mma::mma_sync_m16n16k16_row_col_f16f16f32( + s_frag[mma_q][mma_kv], a_frag[mma_q], b_frag); + } + } else if (std::is_same_v) { + if (mma_d == 0) { + mma::mma_sync_m16n16k16_row_col_f16f16f16( + (uint32_t*)s_frag[mma_q][mma_kv], a_frag[mma_q], b_frag); + } else { + mma::mma_sync_m16n16k16_row_col_f16f16f16((uint32_t*)s_frag[mma_q][mma_kv], + a_frag[mma_q], b_frag); + } + } + } + } + if constexpr (sizeof(typename KTraits::DTypeKV) == 1) { + if (mma_d % 2 == 1) { + *k_smem_offset_r = + k_smem->template advance_offset_by_column<2>(*k_smem_offset_r, mma_d / 2); + } + *k_smem_offset_r -= KTraits::NUM_MMA_KV * 16 * UPCAST_STRIDE_K; + } else { + *k_smem_offset_r = k_smem->template advance_offset_by_column<2>(*k_smem_offset_r, mma_d) - + KTraits::NUM_MMA_KV * 16 * UPCAST_STRIDE_K; + } + } + *q_smem_offset_r -= KTraits::NUM_MMA_D_QK * 2; + *k_smem_offset_r -= KTraits::NUM_MMA_D_QK * sizeof(typename KTraits::DTypeKV); +} + +template +__device__ __forceinline__ void logits_transform( + const Params& params, typename KTraits::AttentionVariant variant, const uint32_t batch_idx, + const uint32_t qo_packed_idx_base, const uint32_t kv_idx_base, const uint32_t qo_len, + const uint32_t kv_len, const uint_fastdiv group_size, + DTypeQKAccum (*s_frag)[KTraits::NUM_MMA_KV][8], const dim3 tid = threadIdx, + const uint32_t kv_head_idx = blockIdx.z) { + const uint32_t lane_idx = tid.x; + uint32_t q[KTraits::NUM_MMA_Q][2], r[KTraits::NUM_MMA_Q][2]; + float logits = 0., logitsTransformed = 0.; + +#pragma unroll + for (uint32_t mma_q = 0; mma_q < KTraits::NUM_MMA_Q; ++mma_q) { +#pragma unroll + for (uint32_t j = 0; j < 2; ++j) { + group_size.divmod(qo_packed_idx_base + mma_q * 16 + lane_idx / 4 + 8 * j, q[mma_q][j], + r[mma_q][j]); + } + } + +#pragma unroll + for (uint32_t mma_q = 0; mma_q < KTraits::NUM_MMA_Q; ++mma_q) { +#pragma unroll + for (uint32_t mma_kv = 0; mma_kv < KTraits::NUM_MMA_KV; ++mma_kv) { +#pragma unroll + for (uint32_t reg_id = 0; reg_id < 8; ++reg_id) { + const uint32_t q_idx = q[mma_q][(reg_id % 4) / 2], kv_idx = kv_idx_base + mma_kv * 16 + + 2 * (lane_idx % 4) + + 8 * (reg_id / 4) + reg_id % 2; + const uint32_t qo_head_idx = kv_head_idx * group_size + r[mma_q][(reg_id % 4) / 2]; + +#ifdef FP16_QK_REDUCTION_SUPPORTED + if constexpr (std::is_same::value) { + logits = std::bit_cast(fp16_ieee_to_fp32_value(s_frag[mma_q][mma_kv][reg_id])); + } else if constexpr (!std::is_same::value) { + logits = s_frag[mma_q][mma_kv][reg_id]; + } +#else + static_assert(!std::is_same::value, + "Set -DFP16_QK_REDUCTION_SUPPORTED and install boost_math " + "then recompile to support fp16 reduction"); + logits = s_frag[mma_q][mma_kv][reg_id]; +#endif + logitsTransformed = variant.LogitsTransform(params, logits, batch_idx, q_idx, kv_idx, + qo_head_idx, kv_head_idx); +#ifdef FP16_QK_REDUCTION_SUPPORTED + if constexpr (std::is_same::value) { + s_frag[mma_q][mma_kv][reg_id] = + std::bit_cast(fp16_ieee_from_fp32_value(logitsTransformed)); + } else if constexpr (!std::is_same::value) { + s_frag[mma_q][mma_kv][reg_id] = logitsTransformed; + } +#else + s_frag[mma_q][mma_kv][reg_id] = logitsTransformed; +#endif + } + } + } +} + +template +__device__ __forceinline__ void logits_mask( + const Params& params, typename KTraits::AttentionVariant variant, const uint32_t batch_idx, + const uint32_t qo_packed_idx_base, const uint32_t kv_idx_base, const uint32_t qo_len, + const uint32_t kv_len, const uint32_t chunk_end, const uint_fastdiv group_size, + typename KTraits::DTypeQKAccum (*s_frag)[KTraits::NUM_MMA_KV][8], const dim3 tid = threadIdx, + const uint32_t kv_head_idx = blockIdx.z) { + const uint32_t lane_idx = tid.x; + constexpr uint32_t NUM_MMA_Q = KTraits::NUM_MMA_Q; + constexpr uint32_t NUM_MMA_KV = KTraits::NUM_MMA_KV; + using DTypeQKAccum = typename KTraits::DTypeQKAccum; + constexpr MaskMode MASK_MODE = KTraits::MASK_MODE; + uint32_t q[NUM_MMA_Q][2], r[NUM_MMA_Q][2]; +#pragma unroll + for (uint32_t mma_q = 0; mma_q < NUM_MMA_Q; ++mma_q) { +#pragma unroll + for (uint32_t j = 0; j < 2; ++j) { + group_size.divmod(qo_packed_idx_base + mma_q * 16 + lane_idx / 4 + 8 * j, q[mma_q][j], + r[mma_q][j]); + } + } + +#pragma unroll + for (uint32_t mma_q = 0; mma_q < NUM_MMA_Q; ++mma_q) { +#pragma unroll + for (uint32_t mma_kv = 0; mma_kv < NUM_MMA_KV; ++mma_kv) { +#pragma unroll + for (uint32_t reg_id = 0; reg_id < 8; ++reg_id) { + const uint32_t q_idx = q[mma_q][(reg_id % 4) / 2], kv_idx = kv_idx_base + mma_kv * 16 + + 2 * (lane_idx % 4) + + 8 * (reg_id / 4) + reg_id % 2; + const uint32_t qo_head_idx = kv_head_idx * group_size + r[mma_q][(reg_id % 4) / 2]; + const bool mask = + (!(MASK_MODE == MaskMode::kCausal || MASK_MODE == MaskMode::kMultiItemScoring + ? (kv_idx + qo_len > kv_len + q_idx || (kv_idx >= chunk_end)) + : kv_idx >= chunk_end)) && + variant.LogitsMask(params, batch_idx, q_idx, kv_idx, qo_head_idx, kv_head_idx); + s_frag[mma_q][mma_kv][reg_id] = + (mask) ? s_frag[mma_q][mma_kv][reg_id] : (KTraits::MaskFillValue); + } + } + } +} + +template +__device__ __forceinline__ void logits_mask_multi_item_scoring( + const Params& params, typename KTraits::AttentionVariant variant, const uint32_t batch_idx, + const uint32_t qo_packed_idx_base, const uint32_t kv_idx_base, const uint32_t qo_len, + const uint32_t kv_len, const uint32_t window_left, const uint32_t chunk_end, + const uint_fastdiv group_size, typename KTraits::DTypeQKAccum (*s_frag)[KTraits::NUM_MMA_KV][8], + // new arguments for compact description of mask + const uint32_t prefix_len, uint16_t* token_pos_in_items, const uint32_t lane_idx = threadIdx.x, + const uint32_t kv_head_idx = blockIdx.z) { + constexpr uint32_t NUM_MMA_Q = KTraits::NUM_MMA_Q; + constexpr uint32_t NUM_MMA_KV = KTraits::NUM_MMA_KV; + using DTypeQKAccum = typename KTraits::DTypeQKAccum; + uint32_t q[NUM_MMA_Q][2], r[NUM_MMA_Q][2]; + +#pragma unroll + for (uint32_t mma_q = 0; mma_q < NUM_MMA_Q; ++mma_q) { +#pragma unroll + for (uint32_t j = 0; j < 2; ++j) { + group_size.divmod(qo_packed_idx_base + mma_q * 16 + lane_idx / 4 + 8 * j, q[mma_q][j], + r[mma_q][j]); + } + } + // prefetching global memory to registers + uint16_t token_pos_in_items_regs[NUM_MMA_Q][(4 / 2)]; +#pragma unroll + for (uint32_t mma_q = 0; mma_q < NUM_MMA_Q; ++mma_q) { +#pragma unroll + for (uint32_t eff_reg_id = 0; eff_reg_id < (4 / 2); ++eff_reg_id) { + const uint32_t q_idx = q[mma_q][eff_reg_id]; + // use __ldca to hint compiler to cache in L1 for further reuse by other tiles + const int idx_in_original_seq = q_idx + kv_len - qo_len; + if (idx_in_original_seq >= prefix_len & idx_in_original_seq < kv_len) { + token_pos_in_items_regs[mma_q][eff_reg_id] = + __ldca(token_pos_in_items + idx_in_original_seq - prefix_len); + } + } + } + +#pragma unroll + for (uint32_t mma_q = 0; mma_q < NUM_MMA_Q; ++mma_q) { +#pragma unroll + for (uint32_t mma_kv = 0; mma_kv < NUM_MMA_KV; ++mma_kv) { +#pragma unroll + for (uint32_t reg_id = 0; reg_id < 8; ++reg_id) { + const uint32_t q_idx = q[mma_q][(reg_id % 4) / 2], kv_idx = kv_idx_base + mma_kv * 16 + + 2 * (lane_idx % 4) + + 8 * (reg_id / 4) + reg_id % 2; + const uint32_t qo_head_idx = kv_head_idx * group_size + r[mma_q][(reg_id % 4) / 2]; + const uint32_t idx_in_original_seq = q_idx + kv_len - qo_len; + const bool out_of_boundary = kv_idx > idx_in_original_seq || (kv_idx >= chunk_end) || + kv_idx + window_left < idx_in_original_seq; + const bool is_prefix = idx_in_original_seq < prefix_len; + if (out_of_boundary || is_prefix) { + s_frag[mma_q][mma_kv][reg_id] = + out_of_boundary ? (KTraits::MaskFillValue) : s_frag[mma_q][mma_kv][reg_id]; + } else { + s_frag[mma_q][mma_kv][reg_id] = + (kv_idx < prefix_len | + (idx_in_original_seq < kv_idx + token_pos_in_items_regs[mma_q][((reg_id % 4) / 2)])) + ? s_frag[mma_q][mma_kv][reg_id] + : (KTraits::MaskFillValue); + } + } + } + } +} + +template +__device__ __forceinline__ void update_mdo_states( + typename KTraits::AttentionVariant variant, + typename KTraits::DTypeQKAccum (*s_frag)[KTraits::NUM_MMA_KV][8], + float (*o_frag)[KTraits::NUM_MMA_D_VO][8], typename KTraits::DTypeQKAccum (*m)[2], + float (*d)[2]) { + using DTypeQKAccum = typename KTraits::DTypeQKAccum; + using AttentionVariant = typename KTraits::AttentionVariant; + constexpr bool use_softmax = AttentionVariant::use_softmax; + + if constexpr (use_softmax) { + const float sm_scale = variant.sm_scale_log2; + if constexpr (std::is_same_v) { +#pragma unroll + for (uint32_t mma_q = 0; mma_q < KTraits::NUM_MMA_Q; ++mma_q) { +#pragma unroll + for (uint32_t j = 0; j < 2; ++j) { + float m_prev = m[mma_q][j]; +#pragma unroll + for (uint32_t mma_kv = 0; mma_kv < KTraits::NUM_MMA_KV; ++mma_kv) { + float m_local = + max(max(s_frag[mma_q][mma_kv][j * 2 + 0], s_frag[mma_q][mma_kv][j * 2 + 1]), + max(s_frag[mma_q][mma_kv][j * 2 + 4], s_frag[mma_q][mma_kv][j * 2 + 5])); + m[mma_q][j] = max(m[mma_q][j], m_local); + } + m[mma_q][j] = max(m[mma_q][j], math::shfl_xor_sync(m[mma_q][j], 0x2)); + m[mma_q][j] = max(m[mma_q][j], math::shfl_xor_sync(m[mma_q][j], 0x1)); + + float o_scale = math::ptx_exp2(m_prev * sm_scale - m[mma_q][j] * sm_scale); + d[mma_q][j] *= o_scale; +#pragma unroll + for (uint32_t mma_d = 0; mma_d < KTraits::NUM_MMA_D_VO; ++mma_d) { + o_frag[mma_q][mma_d][j * 2 + 0] *= o_scale; + o_frag[mma_q][mma_d][j * 2 + 1] *= o_scale; + o_frag[mma_q][mma_d][j * 2 + 4] *= o_scale; + o_frag[mma_q][mma_d][j * 2 + 5] *= o_scale; + } +#pragma unroll + for (uint32_t mma_kv = 0; mma_kv < KTraits::NUM_MMA_KV; ++mma_kv) { + s_frag[mma_q][mma_kv][j * 2 + 0] = math::ptx_exp2( + s_frag[mma_q][mma_kv][j * 2 + 0] * sm_scale - m[mma_q][j] * sm_scale); + s_frag[mma_q][mma_kv][j * 2 + 1] = math::ptx_exp2( + s_frag[mma_q][mma_kv][j * 2 + 1] * sm_scale - m[mma_q][j] * sm_scale); + s_frag[mma_q][mma_kv][j * 2 + 4] = math::ptx_exp2( + s_frag[mma_q][mma_kv][j * 2 + 4] * sm_scale - m[mma_q][j] * sm_scale); + s_frag[mma_q][mma_kv][j * 2 + 5] = math::ptx_exp2( + s_frag[mma_q][mma_kv][j * 2 + 5] * sm_scale - m[mma_q][j] * sm_scale); + } + } + } + } else if constexpr (std::is_same_v) { + const half2 sm_scale = __float2half2_rn(variant.sm_scale_log2); +#pragma unroll + for (uint32_t mma_q = 0; mma_q < KTraits::NUM_MMA_Q; ++mma_q) { + half m_prev[2]; +#pragma unroll + for (uint32_t j = 0; j < 2; ++j) { + m_prev[j] = m[mma_q][j]; +#pragma unroll + for (uint32_t mma_kv = 0; mma_kv < KTraits::NUM_MMA_KV; ++mma_kv) { + half2 m_local = __hmax2(*(half2*)&s_frag[mma_q][mma_kv][j * 2], + *(half2*)&s_frag[mma_q][mma_kv][j * 2 + 4]); + m[mma_q][j] = __hmax(m[mma_q][j], __hmax(m_local.x, m_local.y)); + } + } + *(half2*)&m[mma_q] = + __hmax2(*(half2*)&m[mma_q], math::shfl_xor_sync(*(half2*)&m[mma_q], 0x2)); + *(half2*)&m[mma_q] = + __hmax2(*(half2*)&m[mma_q], math::shfl_xor_sync(*(half2*)&m[mma_q], 0x1)); +#pragma unroll + for (uint32_t j = 0; j < 2; ++j) { + float o_scale = math::ptx_exp2(float(m_prev[j] * sm_scale.x - m[mma_q][j] * sm_scale.x)); + d[mma_q][j] *= o_scale; +#pragma unroll + for (uint32_t mma_d = 0; mma_d < KTraits::NUM_MMA_D_VO; ++mma_d) { + o_frag[mma_q][mma_d][j * 2 + 0] *= o_scale; + o_frag[mma_q][mma_d][j * 2 + 1] *= o_scale; + o_frag[mma_q][mma_d][j * 2 + 4] *= o_scale; + o_frag[mma_q][mma_d][j * 2 + 5] *= o_scale; + } + half2 m2 = make_half2(m[mma_q][j], m[mma_q][j]); +#pragma unroll + for (uint32_t mma_kv = 0; mma_kv < KTraits::NUM_MMA_KV; ++mma_kv) { + *(half2*)&s_frag[mma_q][mma_kv][j * 2] = + math::ptx_exp2(*(half2*)&s_frag[mma_q][mma_kv][j * 2] * sm_scale - m2 * sm_scale); + *(half2*)&s_frag[mma_q][mma_kv][j * 2 + 4] = math::ptx_exp2( + *(half2*)&s_frag[mma_q][mma_kv][j * 2 + 4] * sm_scale - m2 * sm_scale); + } + } + } + } + } +} + +template +__device__ __forceinline__ void compute_sfm_v( + smem_t* v_smem, uint32_t* v_smem_offset_r, + typename KTraits::DTypeQKAccum (*s_frag)[KTraits::NUM_MMA_KV][8], + float (*o_frag)[KTraits::NUM_MMA_D_VO][8], float (*d)[2]) { + constexpr uint32_t UPCAST_STRIDE_V = KTraits::UPCAST_STRIDE_V; + + typename KTraits::DTypeQ s_frag_f16[KTraits::NUM_MMA_Q][KTraits::NUM_MMA_KV][8]; + if constexpr (std::is_same_v) { +#pragma unroll + for (uint32_t mma_q = 0; mma_q < KTraits::NUM_MMA_Q; ++mma_q) { +#pragma unroll + for (uint32_t mma_kv = 0; mma_kv < KTraits::NUM_MMA_KV; ++mma_kv) { + vec_cast::cast<8>(s_frag_f16[mma_q][mma_kv], + s_frag[mma_q][mma_kv]); + } + } + } + + if constexpr (KTraits::AttentionVariant::use_softmax) { +#pragma unroll + for (uint32_t mma_q = 0; mma_q < KTraits::NUM_MMA_Q; ++mma_q) { +#pragma unroll + for (uint32_t mma_kv = 0; mma_kv < KTraits::NUM_MMA_KV; ++mma_kv) { + if constexpr (std::is_same_v) { + mma::m16k16_rowsum_f16f16f32(d[mma_q], s_frag_f16[mma_q][mma_kv]); + } else { + mma::m16k16_rowsum_f16f16f32(d[mma_q], s_frag[mma_q][mma_kv]); + } + } + } + } + +#pragma unroll + for (uint32_t mma_kv = 0; mma_kv < KTraits::NUM_MMA_KV; ++mma_kv) { +#pragma unroll + for (uint32_t mma_d = 0; mma_d < KTraits::NUM_MMA_D_VO; ++mma_d) { + uint32_t b_frag[4]; + if constexpr (sizeof(typename KTraits::DTypeKV) == 1) { + uint32_t b_frag_f8[2]; + if (mma_d % 2 == 0) { + v_smem->ldmatrix_m8n8x4_trans_left_half(*v_smem_offset_r, b_frag_f8); + } else { + v_smem->ldmatrix_m8n8x4_trans_right_half(*v_smem_offset_r, b_frag_f8); + } + b_frag_f8[0] = frag_layout_swizzle_16b_to_8b_trans(b_frag_f8[0]); + b_frag_f8[1] = frag_layout_swizzle_16b_to_8b_trans(b_frag_f8[1]); + vec_cast::cast<8>( + (typename KTraits::DTypeQ*)b_frag, (typename KTraits::DTypeKV*)b_frag_f8); + swap(b_frag[1], b_frag[2]); + } else { + v_smem->ldmatrix_m8n8x4_trans(*v_smem_offset_r, b_frag); + } +#pragma unroll + for (uint32_t mma_q = 0; mma_q < KTraits::NUM_MMA_Q; ++mma_q) { + if constexpr (std::is_same_v) { + mma::mma_sync_m16n16k16_row_col_f16f16f32( + o_frag[mma_q][mma_d], (uint32_t*)s_frag_f16[mma_q][mma_kv], b_frag); + } else { + mma::mma_sync_m16n16k16_row_col_f16f16f32( + o_frag[mma_q][mma_d], (uint32_t*)s_frag[mma_q][mma_kv], b_frag); + } + } + if constexpr (sizeof(typename KTraits::DTypeKV) == 1) { + if (mma_d % 2 == 1) { + *v_smem_offset_r = + v_smem->template advance_offset_by_column<2>(*v_smem_offset_r, mma_d / 2); + } + } else { + *v_smem_offset_r = v_smem->template advance_offset_by_column<2>(*v_smem_offset_r, mma_d); + } + } + *v_smem_offset_r = + v_smem->template advance_offset_by_row<16, UPCAST_STRIDE_V>(*v_smem_offset_r) - + sizeof(typename KTraits::DTypeKV) * KTraits::NUM_MMA_D_VO; + } + *v_smem_offset_r -= 16 * KTraits::NUM_MMA_KV * UPCAST_STRIDE_V; +} + +template +__device__ __forceinline__ void finalize_m(typename KTraits::AttentionVariant variant, + typename KTraits::DTypeQKAccum (*m)[2]) { + if constexpr (variant.use_softmax) { +#pragma unroll + for (uint32_t mma_q = 0; mma_q < KTraits::NUM_MMA_Q; ++mma_q) { +#pragma unroll + for (uint32_t j = 0; j < 2; ++j) { + if (m[mma_q][j] != typename KTraits::DTypeQKAccum(-math::inf)) { + m[mma_q][j] *= variant.sm_scale_log2; + } + } + } + } +} + +template +__device__ __forceinline__ void transform_output( + const Params& params, typename KTraits::AttentionVariant variant, + float (*o_frag)[KTraits::NUM_MMA_D_VO][8], typename KTraits::DTypeQKAccum (*m)[2], + float (*d)[2], const uint32_t batch_idx, const uint32_t kv_tile_idx, + const uint32_t qo_packed_idx_base, const uint32_t warp_idx, const uint32_t lane_idx, + uint32_t kv_head_idx, const uint_fastdiv group_size) { + uint32_t q[KTraits::NUM_MMA_Q][2], r[KTraits::NUM_MMA_Q][2]; + float scale[KTraits::NUM_MMA_Q][2]; +#pragma unroll + for (uint32_t mma_q = 0; mma_q < KTraits::NUM_MMA_Q; ++mma_q) { +#pragma unroll + for (uint32_t j = 0; j < 2; ++j) { + group_size.divmod(qo_packed_idx_base + mma_q * 16 + lane_idx / 4 + 8 * j, q[mma_q][j], + r[mma_q][j]); + uint32_t qo_head_idx = kv_head_idx * group_size + r[mma_q][j]; + // Update the m and d when attention sinks are used. + variant.update_m_d(params, kv_tile_idx, qo_head_idx, m[mma_q][j], d[mma_q][j], + scale[mma_q][j]); + } + } + +#pragma unroll + for (uint32_t mma_q = 0; mma_q < KTraits::NUM_MMA_Q; ++mma_q) { +#pragma unroll + for (uint32_t mma_d = 0; mma_d < KTraits::NUM_MMA_D_VO; ++mma_d) { +#pragma unroll + for (uint32_t reg_id = 0; reg_id < 8; ++reg_id) { + const uint32_t qo_idx = q[mma_q][(reg_id % 4) / 2]; + const uint32_t qo_head_idx = kv_head_idx * group_size + r[mma_q][(reg_id % 4) / 2]; + o_frag[mma_q][mma_d][reg_id] = variant.OutputTransform( + params, o_frag[mma_q][mma_d][reg_id], batch_idx, qo_idx, qo_head_idx, + m[mma_q][(reg_id % 4) / 2], d[mma_q][(reg_id % 4) / 2], scale[mma_q][(reg_id % 4) / 2]); + } + } + } +} + +/*! + * \brief Synchronize the states of the MDO kernel across the threadblock along threadIdx.z. + */ +template +__device__ __forceinline__ void threadblock_sync_mdo_states( + float (*o_frag)[KTraits::NUM_MMA_D_VO][8], typename KTraits::SharedStorage* smem_storage, + typename KTraits::DTypeQKAccum (*m)[2], float (*d)[2], const uint32_t warp_idx, + const uint32_t lane_idx, const dim3 tid = threadIdx) { + // only necessary when blockDim.z > 1 + if constexpr (KTraits::NUM_WARPS_KV > 1) { + float* smem_o = smem_storage->cta_sync_o_smem; + float2* smem_md = smem_storage->cta_sync_md_smem; + // o: [num_warps, NUM_MMA_Q, NUM_MMA_D_VO, WARP_SIZE(32), 8] + // md: [num_warps, NUM_MMA_Q, 16, 2 (m/d)] +#pragma unroll + for (uint32_t mma_q = 0; mma_q < KTraits::NUM_MMA_Q; ++mma_q) { +#pragma unroll + for (uint32_t mma_d = 0; mma_d < KTraits::NUM_MMA_D_VO; ++mma_d) { + vec_t::memcpy( + smem_o + (((warp_idx * KTraits::NUM_MMA_Q + mma_q) * KTraits::NUM_MMA_D_VO + mma_d) * + WARP_SIZE + + lane_idx) * + 8, + o_frag[mma_q][mma_d]); + } + } + + if constexpr (KTraits::AttentionVariant::use_softmax) { +#pragma unroll + for (uint32_t mma_q = 0; mma_q < KTraits::NUM_MMA_Q; ++mma_q) { +#pragma unroll + for (uint32_t j = 0; j < 2; ++j) { + smem_md[((warp_idx * KTraits::NUM_MMA_Q + mma_q) * 2 + j) * 8 + lane_idx / 4] = + make_float2(float(m[mma_q][j]), d[mma_q][j]); + } + } + + // synchronize m,d first + __syncthreads(); +#pragma unroll + for (uint32_t mma_q = 0; mma_q < KTraits::NUM_MMA_Q; ++mma_q) { + float o_scale[2][KTraits::NUM_WARPS_KV]; +#pragma unroll + for (uint32_t j = 0; j < 2; ++j) { + float m_new = -math::inf, d_new = 1.f; +#pragma unroll + for (uint32_t i = 0; i < KTraits::NUM_WARPS_KV; ++i) { + float2 md = smem_md[(((i * KTraits::NUM_WARPS_Q + get_warp_idx_q(tid.y)) * + KTraits::NUM_MMA_Q + + mma_q) * + 2 + + j) * + 8 + + lane_idx / 4]; + float m_prev = m_new, d_prev = d_new; + m_new = max(m_new, md.x); + d_new = d_prev * math::ptx_exp2(m_prev - m_new) + md.y * math::ptx_exp2(md.x - m_new); + } + +#pragma unroll + for (uint32_t i = 0; i < KTraits::NUM_WARPS_KV; ++i) { + float2 md = smem_md[(((i * KTraits::NUM_WARPS_Q + get_warp_idx_q(tid.y)) * + KTraits::NUM_MMA_Q + + mma_q) * + 2 + + j) * + 8 + + lane_idx / 4]; + float mi = md.x; + o_scale[j][i] = math::ptx_exp2(float(mi - m_new)); + } + m[mma_q][j] = typename KTraits::DTypeQKAccum(m_new); + d[mma_q][j] = d_new; + } + +#pragma unroll + for (uint32_t mma_d = 0; mma_d < KTraits::NUM_MMA_D_VO; ++mma_d) { + vec_t o_new; + o_new.fill(0.f); +#pragma unroll + for (uint32_t i = 0; i < KTraits::NUM_WARPS_KV; ++i) { + vec_t oi; + oi.load(smem_o + ((((i * KTraits::NUM_WARPS_Q + get_warp_idx_q(tid.y)) * + KTraits::NUM_MMA_Q + + mma_q) * + KTraits::NUM_MMA_D_VO + + mma_d) * + WARP_SIZE + + lane_idx) * + 8); + +#pragma unroll + for (uint32_t reg_id = 0; reg_id < 8; ++reg_id) { + o_new[reg_id] += oi[reg_id] * o_scale[(reg_id % 4) / 2][i]; + } + } + o_new.store(o_frag[mma_q][mma_d]); + } + } + } else { + // synchronize m,d first + __syncthreads(); +#pragma unroll + for (uint32_t mma_q = 0; mma_q < KTraits::NUM_MMA_Q; ++mma_q) { +#pragma unroll + for (uint32_t mma_d = 0; mma_d < KTraits::NUM_MMA_D_VO; ++mma_d) { + vec_t o_new; + o_new.fill(0.f); +#pragma unroll + for (uint32_t i = 0; i < KTraits::NUM_WARPS_KV; ++i) { + vec_t oi; + oi.load(smem_o + ((((i * KTraits::NUM_WARPS_Q + get_warp_idx_q(tid.y)) * + KTraits::NUM_MMA_Q + + mma_q) * + KTraits::NUM_MMA_D_VO + + mma_d) * + WARP_SIZE + + lane_idx) * + 8); +#pragma unroll + for (uint32_t reg_id = 0; reg_id < 8; ++reg_id) { + o_new[reg_id] += oi[reg_id]; + } + } + o_new.store(o_frag[mma_q][mma_d]); + } + } + } + } +} + +template +__device__ __forceinline__ void write_o_reg_gmem( + float (*o_frag)[KTraits::NUM_MMA_D_VO][8], smem_t* o_smem, + typename KTraits::DTypeO* o_ptr_base, const uint32_t o_packed_idx_base, + const uint32_t qo_upper_bound, const uint32_t o_stride_n, const uint32_t o_stride_h, + const uint_fastdiv group_size, const dim3 tid = threadIdx) { + using DTypeO = typename KTraits::DTypeO; + constexpr uint32_t UPCAST_STRIDE_O = KTraits::UPCAST_STRIDE_O; + const uint32_t warp_idx_x = get_warp_idx_q(tid.y); + const uint32_t lane_idx = tid.x; + + if constexpr (sizeof(DTypeO) == 4) { +#pragma unroll + for (uint32_t mma_q = 0; mma_q < KTraits::NUM_MMA_Q; ++mma_q) { +#pragma unroll + for (uint32_t j = 0; j < 2; ++j) { + uint32_t q, r; + group_size.divmod(o_packed_idx_base + lane_idx / 4 + mma_q * 16 + j * 8, q, r); + const uint32_t o_idx = q; +#pragma unroll + for (uint32_t mma_d = 0; mma_d < KTraits::NUM_MMA_D_VO; ++mma_d) { + if (o_idx < qo_upper_bound) { + *reinterpret_cast(o_ptr_base + q * o_stride_n + r * o_stride_h + mma_d * 16 + + (lane_idx % 4) * 2) = + *reinterpret_cast(&o_frag[mma_q][mma_d][j * 2]); + *reinterpret_cast(o_ptr_base + q * o_stride_n + r * o_stride_h + mma_d * 16 + + 8 + (lane_idx % 4) * 2) = + *reinterpret_cast(&o_frag[mma_q][mma_d][4 + j * 2]); + } + } + } + } + } else { + if (get_warp_idx_kv(tid.z) == 0) { +#pragma unroll + for (uint32_t mma_q = 0; mma_q < KTraits::NUM_MMA_Q; ++mma_q) { +#pragma unroll + for (uint32_t mma_d = 0; mma_d < KTraits::NUM_MMA_D_VO; ++mma_d) { + uint32_t o_frag_f16[8 / 2]; + vec_cast::cast<8>((DTypeO*)o_frag_f16, o_frag[mma_q][mma_d]); + +#ifdef FLASHINFER_STMATRIX_M8N8X4_ENABLED + uint32_t o_smem_offset_w = o_smem->get_permuted_offset( + (warp_idx_x * KTraits::NUM_MMA_Q + mma_q) * 16 + lane_idx % 16, + mma_d * 2 + lane_idx / 16); + o_smem->stmatrix_m8n8x4(o_smem_offset_w, o_frag_f16); +#else + uint32_t o_smem_offset_w = o_smem->get_permuted_offset( + (warp_idx_x * KTraits::NUM_MMA_Q + mma_q) * 16 + lane_idx / 4, mma_d * 2); + ((uint32_t*)(o_smem->base + o_smem_offset_w))[lane_idx % 4] = o_frag_f16[0]; + ((uint32_t*)(o_smem->base + o_smem_offset_w + 8 * UPCAST_STRIDE_O))[lane_idx % 4] = + o_frag_f16[1]; + ((uint32_t*)(o_smem->base + (o_smem_offset_w ^ 0x1)))[lane_idx % 4] = o_frag_f16[2]; + ((uint32_t*)(o_smem->base + (o_smem_offset_w ^ 0x1) + + 8 * UPCAST_STRIDE_O))[lane_idx % 4] = o_frag_f16[3]; +#endif + } + } + + uint32_t o_smem_offset_w = o_smem->get_permuted_offset( + warp_idx_x * KTraits::NUM_MMA_Q * 16 + lane_idx / 8, lane_idx % 8); + +#pragma unroll + for (uint32_t mma_q = 0; mma_q < KTraits::NUM_MMA_Q; ++mma_q) { +#pragma unroll + for (uint32_t j = 0; j < 2 * 2; ++j) { + uint32_t q, r; + group_size.divmod(o_packed_idx_base + lane_idx / 8 + mma_q * 16 + j * 4, q, r); + const uint32_t o_idx = q; + DTypeO* o_ptr = + o_ptr_base + q * o_stride_n + r * o_stride_h + (lane_idx % 8) * upcast_size(); +#pragma unroll + for (uint32_t mma_do = 0; mma_do < KTraits::NUM_MMA_D_VO / 4; ++mma_do) { + if (o_idx < qo_upper_bound) { + o_smem->store_128b(o_smem_offset_w, o_ptr); + } + o_ptr += 8 * upcast_size(); + o_smem_offset_w = o_smem->template advance_offset_by_column<8>(o_smem_offset_w, mma_do); + } + o_smem_offset_w = + o_smem->template advance_offset_by_row<4, UPCAST_STRIDE_O>(o_smem_offset_w) - + 2 * KTraits::NUM_MMA_D_VO; + } + } + } + } +} + +} // namespace + +/*! + * \brief FlashAttention prefill CUDA kernel for a single request. + * \tparam partition_kv Whether to split kv_len into chunks. + * \tparam mask_mode The mask mode used in the attention operation. + * \tparam POS_ENCODING_MODE The positional encoding mode. + * \tparam NUM_MMA_Q The number of fragments in x dimension. + * \tparam NUM_MMA_D_VO The number of fragments in y dimension. + * \tparam NUM_MMA_KV The number of fragments in z dimension. + * \tparam num_warps The number of warps in the threadblock. + * \tparam DTypeQ The data type of the query tensor. + * \tparam DTypeKV The data type of the key/value tensor. + * \tparam DTypeO The data type of the output tensor. + * \param q The query tensor. + * \param k The key tensor. + * \param v The value tensor. + * \param o The output tensor. + * \param tmp The temporary buffer (used when partition_kv is true). + * \param lse The logsumexp value. + * \param rope_rcp_scale 1/(rope_scale), where rope_scale is the scaling + * factor used in RoPE interpolation. + * \param rope_rcp_theta 1/(rope_theta), where rope_theta is the theta + * used in RoPE. + */ +template +__device__ __forceinline__ void SinglePrefillWithKVCacheDevice( + const Params params, typename KTraits::SharedStorage& smem_storage, const dim3 tid = threadIdx, + const uint32_t bx = blockIdx.x, const uint32_t chunk_idx = blockIdx.y, + const uint32_t kv_head_idx = blockIdx.z, const uint32_t num_chunks = gridDim.y, + const uint32_t num_kv_heads = gridDim.z) { + using DTypeQ = typename Params::DTypeQ; +#if (__CUDA_ARCH__ < 800) + if constexpr (std::is_same_v) { + FLASHINFER_RUNTIME_ASSERT("Prefill kernels do not support bf16 on sm75."); + } else { +#endif + using DTypeKV = typename Params::DTypeKV; + using DTypeO = typename Params::DTypeO; + using DTypeQKAccum = typename KTraits::DTypeQKAccum; + using AttentionVariant = typename KTraits::AttentionVariant; + [[maybe_unused]] constexpr uint32_t NUM_MMA_Q = KTraits::NUM_MMA_Q; + [[maybe_unused]] constexpr uint32_t NUM_MMA_KV = KTraits::NUM_MMA_KV; + [[maybe_unused]] constexpr uint32_t NUM_MMA_D_QK = KTraits::NUM_MMA_D_QK; + [[maybe_unused]] constexpr uint32_t NUM_MMA_D_VO = KTraits::NUM_MMA_D_VO; + [[maybe_unused]] constexpr uint32_t HEAD_DIM_QK = KTraits::HEAD_DIM_QK; + [[maybe_unused]] constexpr uint32_t HEAD_DIM_VO = KTraits::HEAD_DIM_VO; + [[maybe_unused]] constexpr uint32_t UPCAST_STRIDE_Q = KTraits::UPCAST_STRIDE_Q; + [[maybe_unused]] constexpr uint32_t UPCAST_STRIDE_K = KTraits::UPCAST_STRIDE_K; + [[maybe_unused]] constexpr uint32_t UPCAST_STRIDE_V = KTraits::UPCAST_STRIDE_V; + [[maybe_unused]] constexpr uint32_t UPCAST_STRIDE_O = KTraits::UPCAST_STRIDE_O; + [[maybe_unused]] constexpr uint32_t CTA_TILE_Q = KTraits::CTA_TILE_Q; + [[maybe_unused]] constexpr uint32_t CTA_TILE_KV = KTraits::CTA_TILE_KV; + [[maybe_unused]] constexpr uint32_t NUM_WARPS_Q = KTraits::NUM_WARPS_Q; + [[maybe_unused]] constexpr uint32_t NUM_WARPS_KV = KTraits::NUM_WARPS_KV; + [[maybe_unused]] constexpr SwizzleMode SWIZZLE_MODE_Q = KTraits::SWIZZLE_MODE_Q; + [[maybe_unused]] constexpr SwizzleMode SWIZZLE_MODE_KV = KTraits::SWIZZLE_MODE_KV; + [[maybe_unused]] constexpr uint32_t KV_THR_LAYOUT_ROW = KTraits::KV_THR_LAYOUT_ROW; + [[maybe_unused]] constexpr uint32_t KV_THR_LAYOUT_COL = KTraits::KV_THR_LAYOUT_COL; + [[maybe_unused]] constexpr MaskMode MASK_MODE = KTraits::MASK_MODE; + + DTypeQ* q = params.q; + DTypeKV* k = params.k; + DTypeKV* v = params.v; + DTypeO* o = params.o; + float* lse = params.lse; + const uint32_t qo_len = params.qo_len; + const uint32_t kv_len = params.kv_len; + const bool partition_kv = params.partition_kv; + const uint32_t q_stride_n = params.q_stride_n; + const uint32_t q_stride_h = params.q_stride_h; + const uint32_t k_stride_n = params.k_stride_n; + const uint32_t k_stride_h = params.k_stride_h; + const uint32_t v_stride_n = params.v_stride_n; + const uint32_t v_stride_h = params.v_stride_h; + const int32_t maybe_window_left = params.window_left; + const uint_fastdiv& group_size = params.group_size; + + static_assert(sizeof(DTypeQ) == 2); + const uint32_t lane_idx = tid.x, warp_idx = get_warp_idx(tid.y, tid.z); + const uint32_t num_qo_heads = num_kv_heads * group_size; + + const uint32_t max_chunk_size = partition_kv ? ceil_div(kv_len, num_chunks) : kv_len; + const uint32_t chunk_start = partition_kv ? chunk_idx * max_chunk_size : 0; + const uint32_t chunk_end = + partition_kv ? min((chunk_idx + 1) * max_chunk_size, kv_len) : kv_len; + const uint32_t chunk_size = chunk_end - chunk_start; + + auto block = cg::this_thread_block(); + auto smem = reinterpret_cast(&smem_storage); + AttentionVariant variant(params, /*batch_idx=*/0, smem); + const uint32_t window_left = variant.window_left; + + DTypeQKAccum s_frag[NUM_MMA_Q][NUM_MMA_KV][8]; + alignas(16) float o_frag[NUM_MMA_Q][NUM_MMA_D_VO][8]; + DTypeQKAccum m[NUM_MMA_Q][2]; + float d[NUM_MMA_Q][2]; + float rope_freq[NUM_MMA_D_QK / 2][4]; + if constexpr (KTraits::POS_ENCODING_MODE == PosEncodingMode::kRoPELlama) { + const float rope_rcp_scale = params.rope_rcp_scale; + const float rope_rcp_theta = params.rope_rcp_theta; + init_rope_freq(rope_freq, rope_rcp_scale, rope_rcp_theta, tid.x); + } + init_states(variant, o_frag, m, d); + + // cooperative fetch q fragment from gmem to reg + const uint32_t qo_packed_idx_base = + (bx * NUM_WARPS_Q + get_warp_idx_q(tid.y)) * NUM_MMA_Q * 16; + smem_t qo_smem(smem_storage.q_smem); + const uint32_t o_stride_n = num_qo_heads * HEAD_DIM_VO, o_stride_h = HEAD_DIM_VO; + DTypeQ* q_ptr_base = q + (kv_head_idx * group_size) * q_stride_h; + DTypeO* o_ptr_base = partition_kv + ? o + chunk_idx * o_stride_n + (kv_head_idx * group_size) * o_stride_h + : o + (kv_head_idx * group_size) * o_stride_h; + + uint32_t q_smem_offset_r = qo_smem.get_permuted_offset( + get_warp_idx_q(tid.y) * NUM_MMA_Q * 16 + lane_idx % 16, lane_idx / 16); + load_q_global_smem(qo_packed_idx_base, qo_len, q_ptr_base, q_stride_n, q_stride_h, + group_size, &qo_smem, tid); + + cp_async::commit_group(); + if constexpr (KTraits::POS_ENCODING_MODE == PosEncodingMode::kRoPELlama) { + cp_async::wait_group<0>(); + block.sync(); + q_smem_inplace_apply_rotary(qo_packed_idx_base, qo_len, kv_len, group_size, &qo_smem, + &q_smem_offset_r, rope_freq, tid); + block.sync(); + } + + smem_t k_smem(smem_storage.k_smem), v_smem(smem_storage.v_smem); + + const uint32_t num_iterations = ceil_div( + MASK_MODE == MaskMode::kCausal + ? min(chunk_size, + sub_if_greater_or_zero( + kv_len - qo_len + ceil_div(((bx + 1) * CTA_TILE_Q), group_size), chunk_start)) + : chunk_size, + CTA_TILE_KV); + + const uint32_t window_iteration = + ceil_div(sub_if_greater_or_zero(kv_len + ceil_div((bx + 1) * CTA_TILE_Q, group_size), + qo_len + window_left + chunk_start), + CTA_TILE_KV); + + const uint32_t mask_iteration = + (MASK_MODE == MaskMode::kCausal + ? min(chunk_size, + sub_if_greater_or_zero(kv_len + ceil_div((bx * CTA_TILE_Q), group_size) - qo_len, + chunk_start)) + : chunk_size) / + CTA_TILE_KV; + + DTypeKV* k_ptr = + k + + (chunk_start + warp_idx * KV_THR_LAYOUT_ROW + lane_idx / KV_THR_LAYOUT_COL) * k_stride_n + + kv_head_idx * k_stride_h + (lane_idx % KV_THR_LAYOUT_COL) * upcast_size(); + DTypeKV* v_ptr = + v + + (chunk_start + warp_idx * KV_THR_LAYOUT_ROW + lane_idx / KV_THR_LAYOUT_COL) * v_stride_n + + kv_head_idx * v_stride_h + (lane_idx % KV_THR_LAYOUT_COL) * upcast_size(); + + uint32_t k_smem_offset_r = k_smem.template get_permuted_offset( + get_warp_idx_kv(tid.z) * NUM_MMA_KV * 16 + 8 * (lane_idx / 16) + + lane_idx % 8, + (lane_idx % 16) / 8), + v_smem_offset_r = v_smem.template get_permuted_offset( + get_warp_idx_kv(tid.z) * NUM_MMA_KV * 16 + lane_idx % 16, lane_idx / 16), + k_smem_offset_w = k_smem.template get_permuted_offset( + warp_idx * KV_THR_LAYOUT_ROW + lane_idx / KV_THR_LAYOUT_COL, + lane_idx % KV_THR_LAYOUT_COL), + v_smem_offset_w = v_smem.template get_permuted_offset( + warp_idx * KV_THR_LAYOUT_ROW + lane_idx / KV_THR_LAYOUT_COL, + lane_idx % KV_THR_LAYOUT_COL); + produce_kv(k_smem, &k_smem_offset_w, &k_ptr, + k_stride_n, 0, chunk_size, tid); + cp_async::commit_group(); + produce_kv(v_smem, &v_smem_offset_w, &v_ptr, + v_stride_n, 0, chunk_size, tid); + cp_async::commit_group(); + +#pragma unroll 1 + for (uint32_t iter = 0; iter < num_iterations; ++iter) { + cp_async::wait_group<1>(); + block.sync(); + + if constexpr (KTraits::POS_ENCODING_MODE == PosEncodingMode::kRoPELlama) { + k_smem_inplace_apply_rotary(chunk_start + iter * CTA_TILE_KV, &k_smem, + &k_smem_offset_r, rope_freq, tid); + block.sync(); + } + + // compute attention score + compute_qk(&qo_smem, &q_smem_offset_r, &k_smem, &k_smem_offset_r, s_frag); + + logits_transform( + params, variant, /*batch_idx=*/0, qo_packed_idx_base, + chunk_start + (iter * NUM_WARPS_KV + get_warp_idx_kv(tid.z)) * NUM_MMA_KV * 16, + qo_len, kv_len, group_size, s_frag, tid, kv_head_idx); + + // apply mask + if (MASK_MODE == MaskMode::kCustom || (iter >= mask_iteration || iter < window_iteration)) { + logits_mask( + params, variant, /*batch_idx=*/0, qo_packed_idx_base, + chunk_start + (iter * NUM_WARPS_KV + get_warp_idx_kv(tid.z)) * NUM_MMA_KV * 16, + qo_len, kv_len, chunk_end, group_size, s_frag, tid, kv_head_idx); + } + + // compute m,d states in online softmax + update_mdo_states(variant, s_frag, o_frag, m, d); + + block.sync(); + produce_kv( + k_smem, &k_smem_offset_w, &k_ptr, k_stride_n, (iter + 1) * CTA_TILE_KV, chunk_size, tid); + cp_async::commit_group(); + cp_async::wait_group<1>(); + block.sync(); + + // compute sfm*v + compute_sfm_v(&v_smem, &v_smem_offset_r, s_frag, o_frag, d); + + block.sync(); + produce_kv( + v_smem, &v_smem_offset_w, &v_ptr, v_stride_n, (iter + 1) * CTA_TILE_KV, chunk_size, tid); + cp_async::commit_group(); + } + cp_async::wait_group<0>(); + block.sync(); + + finalize_m(variant, m); + + // threadblock synchronization + threadblock_sync_mdo_states(o_frag, &smem_storage, m, d, warp_idx, lane_idx, tid); + + // transform output + transform_output(params, variant, o_frag, m, d, /*batch_idx=*/0, chunk_idx, + qo_packed_idx_base, warp_idx, lane_idx, kv_head_idx, + group_size); + + // write back + write_o_reg_gmem(o_frag, &qo_smem, o_ptr_base, qo_packed_idx_base, qo_len, + /*o_stride_n=*/ + partition_kv ? num_chunks * o_stride_n : o_stride_n, + /*o_stride_h=*/o_stride_h, group_size, tid); + + // write lse + if constexpr (variant.use_softmax) { + if (lse != nullptr || partition_kv) { + if (get_warp_idx_kv(tid.z) == 0) { +#pragma unroll + for (uint32_t mma_q = 0; mma_q < NUM_MMA_Q; ++mma_q) { +#pragma unroll + for (uint32_t j = 0; j < 2; ++j) { + uint32_t q, r; + group_size.divmod(qo_packed_idx_base + lane_idx / 4 + j * 8 + mma_q * 16, q, r); + const uint32_t qo_head_idx = kv_head_idx * group_size + r; + const uint32_t qo_idx = q; + if (qo_idx < qo_len) { + if (partition_kv) { + lse[(qo_idx * num_chunks + chunk_idx) * num_qo_heads + qo_head_idx] = + math::ptx_log2(d[mma_q][j]) + float(m[mma_q][j]); + } else { + lse[qo_idx * num_qo_heads + qo_head_idx] = + math::ptx_log2(d[mma_q][j]) + float(m[mma_q][j]); + } + } + } + } + } + } + } +#if (__CUDA_ARCH__ < 800) + } +#endif +} + +template +__global__ __launch_bounds__(KTraits::NUM_THREADS) void SinglePrefillWithKVCacheKernel( + const __grid_constant__ Params params) { + extern __shared__ uint8_t smem[]; + auto& smem_storage = reinterpret_cast(smem); + SinglePrefillWithKVCacheDevice(params, smem_storage); +} + +template +cudaError_t SinglePrefillWithKVCacheDispatched(Params params, typename Params::DTypeO* tmp, + cudaStream_t stream) { + using DTypeQ = typename Params::DTypeQ; + using DTypeKV = typename Params::DTypeKV; + using DTypeO = typename Params::DTypeO; + const uint32_t num_qo_heads = params.num_qo_heads; + const uint32_t num_kv_heads = params.num_kv_heads; + const uint32_t qo_len = params.qo_len; + const uint32_t kv_len = params.kv_len; + if (kv_len < qo_len && MASK_MODE == MaskMode::kCausal) { + std::ostringstream err_msg; + err_msg << "When mask_mode is set to MaskMode::kCausal, kv_len must be greater than or equal " + "to qo_len, got kv_len" + << kv_len << " and qo_len " << qo_len; + FLASHINFER_ERROR(err_msg.str()); + } + + const uint32_t group_size = num_qo_heads / num_kv_heads; + constexpr uint32_t NUM_MMA_D_QK = HEAD_DIM_QK / 16; + constexpr uint32_t NUM_MMA_D_VO = HEAD_DIM_VO / 16; + int64_t packed_qo_len = qo_len * group_size; + uint32_t cta_tile_q = FA2DetermineCtaTileQ(packed_qo_len, HEAD_DIM_VO); + + DISPATCH_CTA_TILE_Q(cta_tile_q, CTA_TILE_Q, { + constexpr uint32_t NUM_WARPS_Q = get_num_warps_q(CTA_TILE_Q); + constexpr uint32_t NUM_WARPS_KV = get_num_warps_kv(CTA_TILE_Q); + constexpr uint32_t NUM_MMA_Q = get_num_mma_q(CTA_TILE_Q); + + using DTypeQKAccum = + typename std::conditional, half, + float>::type; + + int dev_id = 0; + FLASHINFER_CUDA_CALL(cudaGetDevice(&dev_id)); + int max_smem_per_sm = 0; + FLASHINFER_CUDA_CALL(cudaDeviceGetAttribute( + &max_smem_per_sm, cudaDevAttrMaxSharedMemoryPerMultiprocessor, dev_id)); + // we expect each sm execute two threadblocks + const int num_ctas_per_sm = + max_smem_per_sm >= 2 * (CTA_TILE_Q * HEAD_DIM_QK * sizeof(DTypeQ) + + (HEAD_DIM_QK + HEAD_DIM_VO) * 16 * NUM_WARPS_KV * sizeof(DTypeKV)) + ? 2 + : 1; + const int max_smem_per_threadblock = max_smem_per_sm / num_ctas_per_sm; + + const uint32_t max_num_mma_kv_reg = + (HEAD_DIM_VO >= 128 && NUM_MMA_Q == 2 && POS_ENCODING_MODE == PosEncodingMode::kRoPELlama && + !USE_FP16_QK_REDUCTION) + ? 2 + : (8 / NUM_MMA_Q); + const uint32_t max_num_mma_kv_smem = + (max_smem_per_threadblock - CTA_TILE_Q * HEAD_DIM_QK * sizeof(DTypeQ)) / + ((HEAD_DIM_QK + HEAD_DIM_VO) * 16 * NUM_WARPS_KV * sizeof(DTypeKV)); + + // control NUM_MMA_KV for maximum warp occupancy + DISPATCH_NUM_MMA_KV(min(max_num_mma_kv_smem, max_num_mma_kv_reg), NUM_MMA_KV, { + using KTraits = + KernelTraits; + if constexpr (KTraits::IsInvalid()) { + // Invalid configuration, skip + std::ostringstream err_msg; + err_msg << "FlashInfer Internal Error: Invalid configuration : NUM_MMA_Q=" << NUM_MMA_Q + << " NUM_MMA_D_QK=" << NUM_MMA_D_QK << " NUM_MMA_D_VO=" << NUM_MMA_D_VO + << " NUM_MMA_KV=" << NUM_MMA_KV << " NUM_WARPS_Q=" << NUM_WARPS_Q + << " NUM_WARPS_KV=" << NUM_WARPS_KV + << " please create an issue (https://github.com/flashinfer-ai/flashinfer/issues)" + " and report the issue to the developers."; + FLASHINFER_ERROR(err_msg.str()); + } else { + constexpr uint32_t num_threads = (NUM_WARPS_Q * NUM_WARPS_KV) * WARP_SIZE; + auto kernel = SinglePrefillWithKVCacheKernel; + size_t smem_size = sizeof(typename KTraits::SharedStorage); + FLASHINFER_CUDA_CALL( + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + int num_blocks_per_sm = 0; + int num_sm = 0; + FLASHINFER_CUDA_CALL( + cudaDeviceGetAttribute(&num_sm, cudaDevAttrMultiProcessorCount, dev_id)); + FLASHINFER_CUDA_CALL(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &num_blocks_per_sm, kernel, num_threads, smem_size)); + uint32_t max_num_kv_chunks = (num_blocks_per_sm * num_sm) / + (num_kv_heads * ceil_div(qo_len * group_size, CTA_TILE_Q)); + uint32_t num_chunks; + if (max_num_kv_chunks > 0) { + uint32_t chunk_size = max(ceil_div(kv_len, max_num_kv_chunks), 256); + num_chunks = ceil_div(kv_len, chunk_size); + } else { + num_chunks = 0; + } + + if (num_chunks <= 1 || tmp == nullptr) { + // Enough parallelism, do not split-kv + params.partition_kv = false; + void* args[] = {(void*)¶ms}; + dim3 nblks(ceil_div(qo_len * group_size, CTA_TILE_Q), 1, num_kv_heads); + dim3 nthrs(32, NUM_WARPS_Q, NUM_WARPS_KV); + FLASHINFER_CUDA_CALL( + cudaLaunchKernel((void*)kernel, nblks, nthrs, args, smem_size, stream)); + } else { + // Use cooperative groups to increase occupancy + params.partition_kv = true; + float* tmp_lse = (float*)(tmp + num_chunks * qo_len * num_qo_heads * HEAD_DIM_VO); + auto o = params.o; + auto lse = params.lse; + params.o = tmp; + params.lse = tmp_lse; + void* args[] = {(void*)¶ms}; + dim3 nblks(ceil_div(qo_len * group_size, CTA_TILE_Q), num_chunks, num_kv_heads); + dim3 nthrs(32, NUM_WARPS_Q, NUM_WARPS_KV); + + FLASHINFER_CUDA_CALL( + cudaLaunchKernel((void*)kernel, nblks, nthrs, args, smem_size, stream)); + if constexpr (AttentionVariant::use_softmax) { + FLASHINFER_CUDA_CALL(MergeStates(tmp, tmp_lse, o, lse, num_chunks, qo_len, num_qo_heads, + HEAD_DIM_VO, stream)); + } else { + FLASHINFER_CUDA_CALL( + AttentionSum(tmp, o, num_chunks, qo_len, num_qo_heads, HEAD_DIM_VO, stream)); + } + } + } + }) + }); + return cudaSuccess; +} + +template +__global__ __launch_bounds__(KTraits::NUM_THREADS) void BatchPrefillWithRaggedKVCacheKernel( + const __grid_constant__ Params params) { + using DTypeQ = typename Params::DTypeQ; +#if (__CUDA_ARCH__ < 800) + if constexpr (std::is_same_v) { + FLASHINFER_RUNTIME_ASSERT("Prefill kernels do not support bf16 on sm75."); + } else { +#endif + using DTypeKV = typename Params::DTypeKV; + using DTypeO = typename Params::DTypeO; + using IdType = typename Params::IdType; + using DTypeQKAccum = typename KTraits::DTypeQKAccum; + using AttentionVariant = typename KTraits::AttentionVariant; + [[maybe_unused]] constexpr uint32_t NUM_MMA_Q = KTraits::NUM_MMA_Q; + [[maybe_unused]] constexpr uint32_t NUM_MMA_KV = KTraits::NUM_MMA_KV; + [[maybe_unused]] constexpr uint32_t NUM_MMA_D_QK = KTraits::NUM_MMA_D_QK; + [[maybe_unused]] constexpr uint32_t NUM_MMA_D_VO = KTraits::NUM_MMA_D_VO; + [[maybe_unused]] constexpr uint32_t HEAD_DIM_QK = KTraits::HEAD_DIM_QK; + [[maybe_unused]] constexpr uint32_t HEAD_DIM_VO = KTraits::HEAD_DIM_VO; + [[maybe_unused]] constexpr uint32_t UPCAST_STRIDE_Q = KTraits::UPCAST_STRIDE_Q; + [[maybe_unused]] constexpr uint32_t UPCAST_STRIDE_K = KTraits::UPCAST_STRIDE_K; + [[maybe_unused]] constexpr uint32_t UPCAST_STRIDE_V = KTraits::UPCAST_STRIDE_V; + [[maybe_unused]] constexpr uint32_t UPCAST_STRIDE_O = KTraits::UPCAST_STRIDE_O; + [[maybe_unused]] constexpr uint32_t CTA_TILE_Q = KTraits::CTA_TILE_Q; + [[maybe_unused]] constexpr uint32_t CTA_TILE_KV = KTraits::CTA_TILE_KV; + [[maybe_unused]] constexpr uint32_t NUM_WARPS_Q = KTraits::NUM_WARPS_Q; + [[maybe_unused]] constexpr uint32_t NUM_WARPS_KV = KTraits::NUM_WARPS_KV; + [[maybe_unused]] constexpr SwizzleMode SWIZZLE_MODE_Q = KTraits::SWIZZLE_MODE_Q; + [[maybe_unused]] constexpr SwizzleMode SWIZZLE_MODE_KV = KTraits::SWIZZLE_MODE_KV; + [[maybe_unused]] constexpr uint32_t KV_THR_LAYOUT_ROW = KTraits::KV_THR_LAYOUT_ROW; + [[maybe_unused]] constexpr uint32_t KV_THR_LAYOUT_COL = KTraits::KV_THR_LAYOUT_COL; + [[maybe_unused]] constexpr MaskMode MASK_MODE = KTraits::MASK_MODE; + + DTypeQ* q = params.q; + IdType* request_indices = params.request_indices; + IdType* qo_tile_indices = params.qo_tile_indices; + IdType* kv_tile_indices = params.kv_tile_indices; + IdType* q_indptr = params.q_indptr; + IdType* kv_indptr = params.kv_indptr; + DTypeKV* k = params.k; + DTypeKV* v = params.v; + IdType* o_indptr = params.o_indptr; + DTypeO* o = params.o; + float* lse = params.lse; + bool* block_valid_mask = params.block_valid_mask; + const bool partition_kv = params.partition_kv; + const uint32_t q_stride_n = params.q_stride_n; + const uint32_t q_stride_h = params.q_stride_h; + const uint32_t k_stride_n = params.k_stride_n; + const uint32_t k_stride_h = params.k_stride_h; + const uint32_t v_stride_n = params.v_stride_n; + const uint32_t v_stride_h = params.v_stride_h; + const int32_t maybe_window_left = params.window_left; + const uint_fastdiv& group_size = params.group_size; + + static_assert(sizeof(DTypeQ) == 2); + const uint32_t kv_chunk_size = *(params.kv_chunk_size_ptr); + const dim3& tid = threadIdx; + + auto block = cg::this_thread_block(); + const uint32_t bx = blockIdx.x, lane_idx = tid.x, + warp_idx = get_warp_idx(tid.y, tid.z), kv_head_idx = blockIdx.z; + if (block_valid_mask && !block_valid_mask[bx]) { + return; + } + const uint32_t num_kv_heads = gridDim.z, num_qo_heads = group_size * num_kv_heads; + const uint32_t request_idx = request_indices[bx], qo_tile_idx = qo_tile_indices[bx], + kv_tile_idx = kv_tile_indices[bx]; + extern __shared__ uint8_t smem[]; + auto& smem_storage = reinterpret_cast(smem); + AttentionVariant variant(params, /*batch_idx=*/request_idx, smem); + const uint32_t qo_len = variant.qo_len, kv_len = variant.kv_len, + window_left = variant.window_left; + const uint32_t kv_len_safe = kv_len > 0 ? kv_len : 1; + const uint32_t max_chunk_size = partition_kv ? kv_chunk_size : kv_len; + const uint32_t chunk_start = partition_kv ? kv_tile_idx * max_chunk_size : 0; + const uint32_t chunk_end = + partition_kv ? min((kv_tile_idx + 1) * max_chunk_size, kv_len) : kv_len; + const uint32_t chunk_size = chunk_end - chunk_start; + const uint32_t qo_upper_bound = + min(qo_len, ceil_div((qo_tile_idx + 1) * CTA_TILE_Q, group_size)); + + DTypeQKAccum s_frag[NUM_MMA_Q][NUM_MMA_KV][8]; + alignas(16) float o_frag[NUM_MMA_Q][NUM_MMA_D_VO][8]; + DTypeQKAccum m[NUM_MMA_Q][2]; + float d[NUM_MMA_Q][2]; + float rope_freq[NUM_MMA_D_QK / 2][4]; + + if constexpr (KTraits::POS_ENCODING_MODE == PosEncodingMode::kRoPELlama) { + const float rope_rcp_scale = params.rope_rcp_scale; + const float rope_rcp_theta = params.rope_rcp_theta; + init_rope_freq(rope_freq, rope_rcp_scale, rope_rcp_theta, tid.x); + } + init_states(variant, o_frag, m, d); + + const uint32_t qo_packed_idx_base = + (qo_tile_idx * NUM_WARPS_Q + get_warp_idx_q(tid.y)) * NUM_MMA_Q * 16; + smem_t qo_smem(smem_storage.q_smem); + const uint32_t o_stride_n = num_qo_heads * HEAD_DIM_VO, o_stride_h = HEAD_DIM_VO; + + DTypeQ* q_ptr_base = + q + q_indptr[request_idx] * q_stride_n + kv_head_idx * group_size * q_stride_h; + + DTypeO* o_ptr_base = partition_kv ? o + (o_indptr[request_idx] + kv_tile_idx) * o_stride_n + + (kv_head_idx * group_size) * o_stride_h + : o + o_indptr[request_idx] * o_stride_n + + (kv_head_idx * group_size) * o_stride_h; + + uint32_t q_smem_offset_r = qo_smem.get_permuted_offset( + get_warp_idx_q(tid.y) * NUM_MMA_Q * 16 + lane_idx % 16, lane_idx / 16); + +#if (__CUDACC_VER_MAJOR__ >= 12 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.wait;"); +#endif + + load_q_global_smem(qo_packed_idx_base, qo_upper_bound, q_ptr_base, q_stride_n, + q_stride_h, group_size, &qo_smem, tid); + + cp_async::commit_group(); + + if constexpr (KTraits::POS_ENCODING_MODE == PosEncodingMode::kRoPELlama) { + cp_async::wait_group<0>(); + block.sync(); + IdType* q_rope_offset = nullptr; + + if constexpr (has_maybe_q_rope_offset_v) { + q_rope_offset = params.maybe_q_rope_offset; + } + if (!q_rope_offset) { + q_smem_inplace_apply_rotary(qo_packed_idx_base, qo_len, kv_len, group_size, + &qo_smem, &q_smem_offset_r, rope_freq, tid); + } else { + q_smem_inplace_apply_rotary_with_pos( + qo_packed_idx_base, q_rope_offset + q_indptr[request_idx], &qo_smem, group_size, + &q_smem_offset_r, rope_freq, tid); + } + block.sync(); + } + + const uint32_t num_iterations = ceil_div( + (MASK_MODE == MaskMode::kCausal + ? min(chunk_size, + sub_if_greater_or_zero( + kv_len - qo_len + ceil_div(((qo_tile_idx + 1) * CTA_TILE_Q), group_size), + chunk_start)) + : chunk_size), + CTA_TILE_KV); + + const uint32_t window_iteration = ceil_div( + sub_if_greater_or_zero(kv_len + ceil_div((qo_tile_idx + 1) * CTA_TILE_Q, group_size), + qo_len + window_left + chunk_start), + CTA_TILE_KV); + + const uint32_t mask_iteration = + (MASK_MODE == MaskMode::kCausal + ? min(chunk_size, + sub_if_greater_or_zero( + kv_len + ceil_div((qo_tile_idx * CTA_TILE_Q), group_size) - qo_len, + chunk_start)) + : chunk_size) / + CTA_TILE_KV; + + smem_t k_smem(smem_storage.k_smem), v_smem(smem_storage.v_smem); + + uint32_t k_smem_offset_r = k_smem.template get_permuted_offset( + get_warp_idx_kv(tid.z) * NUM_MMA_KV * 16 + 8 * (lane_idx / 16) + + lane_idx % 8, + (lane_idx % 16) / 8), + v_smem_offset_r = v_smem.template get_permuted_offset( + get_warp_idx_kv(tid.z) * NUM_MMA_KV * 16 + lane_idx % 16, lane_idx / 16), + k_smem_offset_w = k_smem.template get_permuted_offset( + warp_idx * KV_THR_LAYOUT_ROW + lane_idx / KV_THR_LAYOUT_COL, + lane_idx % KV_THR_LAYOUT_COL), + v_smem_offset_w = v_smem.template get_permuted_offset( + warp_idx * KV_THR_LAYOUT_ROW + lane_idx / KV_THR_LAYOUT_COL, + lane_idx % KV_THR_LAYOUT_COL); + + DTypeKV* k_ptr = k + + (kv_indptr[request_idx] + chunk_start + warp_idx * KV_THR_LAYOUT_ROW + + lane_idx / KV_THR_LAYOUT_COL) * + k_stride_n + + kv_head_idx * k_stride_h + + (lane_idx % KV_THR_LAYOUT_COL) * upcast_size(); + DTypeKV* v_ptr = v + + (kv_indptr[request_idx] + chunk_start + warp_idx * KV_THR_LAYOUT_ROW + + lane_idx / KV_THR_LAYOUT_COL) * + v_stride_n + + kv_head_idx * v_stride_h + + (lane_idx % KV_THR_LAYOUT_COL) * upcast_size(); + + produce_kv(k_smem, &k_smem_offset_w, &k_ptr, + k_stride_n, 0, chunk_size, tid); + cp_async::commit_group(); + produce_kv(v_smem, &v_smem_offset_w, &v_ptr, + v_stride_n, 0, chunk_size, tid); + cp_async::commit_group(); + +#pragma unroll 1 + for (uint32_t iter = 0; iter < num_iterations; ++iter) { + cp_async::wait_group<1>(); + block.sync(); + + if constexpr (KTraits::POS_ENCODING_MODE == PosEncodingMode::kRoPELlama) { + IdType* k_rope_offset = nullptr; + if constexpr (has_maybe_k_rope_offset_v) { + k_rope_offset = params.maybe_k_rope_offset; + } + k_smem_inplace_apply_rotary( + (k_rope_offset == nullptr ? 0 : k_rope_offset[request_idx]) + chunk_start + + iter * CTA_TILE_KV, + &k_smem, &k_smem_offset_r, rope_freq, tid); + block.sync(); + } + + // compute attention score + compute_qk(&qo_smem, &q_smem_offset_r, &k_smem, &k_smem_offset_r, s_frag); + + logits_transform( + params, variant, /*batch_idx=*/request_idx, qo_packed_idx_base, + chunk_start + (iter * NUM_WARPS_KV + get_warp_idx_kv(tid.z)) * NUM_MMA_KV * 16, + qo_len, kv_len, group_size, s_frag, tid, kv_head_idx); + + // apply mask + if (MASK_MODE == MaskMode::kCustom || (iter >= mask_iteration || iter < window_iteration)) { + logits_mask( + params, variant, /*batch_idx=*/request_idx, qo_packed_idx_base, + chunk_start + (iter * NUM_WARPS_KV + get_warp_idx_kv(tid.z)) * NUM_MMA_KV * 16, + qo_len, kv_len, chunk_end, group_size, s_frag, tid, kv_head_idx); + } + + // compute m,d states in online softmax + update_mdo_states(variant, s_frag, o_frag, m, d); + + block.sync(); + produce_kv( + k_smem, &k_smem_offset_w, &k_ptr, k_stride_n, (iter + 1) * CTA_TILE_KV, chunk_size, tid); + cp_async::commit_group(); + cp_async::wait_group<1>(); + block.sync(); + + // compute sfm*v + compute_sfm_v(&v_smem, &v_smem_offset_r, s_frag, o_frag, d); + + block.sync(); + produce_kv( + v_smem, &v_smem_offset_w, &v_ptr, v_stride_n, (iter + 1) * CTA_TILE_KV, chunk_size, tid); + cp_async::commit_group(); + } + cp_async::wait_group<0>(); + block.sync(); + + finalize_m(variant, m); + + // threadblock synchronization + threadblock_sync_mdo_states(o_frag, &smem_storage, m, d, warp_idx, lane_idx, tid); + + const uint32_t num_kv_chunks = (kv_len_safe + kv_chunk_size - 1) / kv_chunk_size; + + // transform output + transform_output(params, variant, o_frag, m, d, /*batch_idx=*/request_idx, + kv_tile_idx, qo_packed_idx_base, warp_idx, lane_idx, + kv_head_idx, group_size); + + // write back + write_o_reg_gmem(o_frag, &qo_smem, o_ptr_base, qo_packed_idx_base, qo_len, + /*o_stride_n=*/ + partition_kv ? num_kv_chunks * o_stride_n : o_stride_n, + /*o_stride_h=*/o_stride_h, group_size, tid); + + // write lse + if constexpr (AttentionVariant::use_softmax) { + if (lse != nullptr) { + if (get_warp_idx_kv(tid.z) == 0) { +#pragma unroll + for (uint32_t mma_q = 0; mma_q < NUM_MMA_Q; ++mma_q) { +#pragma unroll + for (uint32_t j = 0; j < 2; ++j) { + uint32_t q, r; + group_size.divmod(qo_packed_idx_base + lane_idx / 4 + j * 8 + mma_q * 16, q, r); + const uint32_t qo_head_idx = kv_head_idx * group_size + r; + const uint32_t qo_idx = q; + if (qo_idx < qo_len) { + if (partition_kv) { + lse[(o_indptr[request_idx] + qo_idx * num_kv_chunks + kv_tile_idx) * + num_qo_heads + + qo_head_idx] = math::ptx_log2(d[mma_q][j]) + float(m[mma_q][j]); + } else { + lse[(o_indptr[request_idx] + qo_idx) * num_qo_heads + qo_head_idx] = + math::ptx_log2(d[mma_q][j]) + float(m[mma_q][j]); + } + } + } + } + } + } + } +#if (__CUDACC_VER_MAJOR__ >= 12 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.launch_dependents;"); +#endif +#if (__CUDA_ARCH__ < 800) + } +#endif +} + +template +__device__ __forceinline__ void BatchPrefillWithPagedKVCacheDevice( + const Params params, typename KTraits::SharedStorage& smem_storage, const dim3 tid = threadIdx, + const uint32_t bx = blockIdx.x, const uint32_t kv_head_idx = blockIdx.z, + const uint32_t num_kv_heads = gridDim.z) { + using DTypeQ = typename Params::DTypeQ; +#if (__CUDA_ARCH__ < 800) + if constexpr (std::is_same_v) { + FLASHINFER_RUNTIME_ASSERT("Prefill kernels do not support bf16 on sm75."); + } else { +#endif + using DTypeKV = typename Params::DTypeKV; + using DTypeO = typename Params::DTypeO; + using IdType = typename Params::IdType; + using DTypeQKAccum = typename KTraits::DTypeQKAccum; + using AttentionVariant = typename KTraits::AttentionVariant; + [[maybe_unused]] constexpr uint32_t NUM_MMA_Q = KTraits::NUM_MMA_Q; + [[maybe_unused]] constexpr uint32_t NUM_MMA_KV = KTraits::NUM_MMA_KV; + [[maybe_unused]] constexpr uint32_t NUM_MMA_D_QK = KTraits::NUM_MMA_D_QK; + [[maybe_unused]] constexpr uint32_t NUM_MMA_D_VO = KTraits::NUM_MMA_D_VO; + [[maybe_unused]] constexpr uint32_t HEAD_DIM_QK = KTraits::HEAD_DIM_QK; + [[maybe_unused]] constexpr uint32_t HEAD_DIM_VO = KTraits::HEAD_DIM_VO; + [[maybe_unused]] constexpr uint32_t UPCAST_STRIDE_Q = KTraits::UPCAST_STRIDE_Q; + [[maybe_unused]] constexpr uint32_t UPCAST_STRIDE_K = KTraits::UPCAST_STRIDE_K; + [[maybe_unused]] constexpr uint32_t UPCAST_STRIDE_V = KTraits::UPCAST_STRIDE_V; + [[maybe_unused]] constexpr uint32_t UPCAST_STRIDE_O = KTraits::UPCAST_STRIDE_O; + [[maybe_unused]] constexpr uint32_t NUM_WARPS_Q = KTraits::NUM_WARPS_Q; + [[maybe_unused]] constexpr uint32_t NUM_WARPS_KV = KTraits::NUM_WARPS_KV; + [[maybe_unused]] constexpr SwizzleMode SWIZZLE_MODE_Q = KTraits::SWIZZLE_MODE_Q; + [[maybe_unused]] constexpr SwizzleMode SWIZZLE_MODE_KV = KTraits::SWIZZLE_MODE_KV; + [[maybe_unused]] constexpr uint32_t CTA_TILE_Q = KTraits::CTA_TILE_Q; + [[maybe_unused]] constexpr uint32_t CTA_TILE_KV = KTraits::CTA_TILE_KV; + [[maybe_unused]] constexpr uint32_t KV_THR_LAYOUT_ROW = KTraits::KV_THR_LAYOUT_ROW; + [[maybe_unused]] constexpr uint32_t KV_THR_LAYOUT_COL = KTraits::KV_THR_LAYOUT_COL; + [[maybe_unused]] constexpr MaskMode MASK_MODE = KTraits::MASK_MODE; + + IdType* request_indices = params.request_indices; + IdType* qo_tile_indices = params.qo_tile_indices; + IdType* kv_tile_indices = params.kv_tile_indices; + DTypeQ* q = params.q; + IdType* q_indptr = params.q_indptr; + IdType* o_indptr = params.o_indptr; + DTypeO* o = params.o; + float* lse = params.lse; + bool* block_valid_mask = params.block_valid_mask; + const paged_kv_t& paged_kv = params.paged_kv; + const bool partition_kv = params.partition_kv; + const int32_t maybe_window_left = params.window_left; + const uint_fastdiv& group_size = params.group_size; + + uint32_t* maybe_prefix_len_ptr = nullptr; + if constexpr (has_maybe_prefix_len_ptr_v) { + maybe_prefix_len_ptr = params.maybe_prefix_len_ptr; + } + uint16_t* maybe_token_pos_in_items_ptr = nullptr; + if constexpr (has_maybe_token_pos_in_items_ptr_v) { + maybe_token_pos_in_items_ptr = params.maybe_token_pos_in_items_ptr; + } + uint32_t token_pos_in_items_len = 0; + if constexpr (has_token_pos_in_items_len_v) { + token_pos_in_items_len = params.token_pos_in_items_len; + } + uint16_t* maybe_max_item_len_ptr = nullptr; + if constexpr (has_maybe_max_item_len_ptr_v) { + maybe_max_item_len_ptr = params.maybe_max_item_len_ptr; + } + + static_assert(sizeof(DTypeQ) == 2); + auto block = cg::this_thread_block(); + const uint32_t kv_chunk_size = *(params.kv_chunk_size_ptr); + + const uint32_t lane_idx = tid.x, warp_idx = get_warp_idx(tid.y, tid.z); + if (block_valid_mask && !block_valid_mask[bx]) { + return; + } + const uint32_t num_qo_heads = num_kv_heads * group_size; + + const uint32_t request_idx = request_indices[bx], qo_tile_idx = qo_tile_indices[bx], + kv_tile_idx = kv_tile_indices[bx]; + auto smem = reinterpret_cast(&smem_storage); + AttentionVariant variant(params, /*batch_idx=*/request_idx, smem); + const uint32_t qo_len = variant.qo_len, kv_len = variant.kv_len, + window_left = variant.window_left; + const uint32_t kv_len_safe = kv_len > 0 ? kv_len : 1; + const uint32_t max_chunk_size = partition_kv ? kv_chunk_size : kv_len; + const uint32_t chunk_start = partition_kv ? kv_tile_idx * max_chunk_size : 0; + const uint32_t chunk_end = + partition_kv ? min((kv_tile_idx + 1) * max_chunk_size, kv_len) : kv_len; + const uint32_t chunk_size = chunk_end - chunk_start; + const uint32_t qo_upper_bound = + min(qo_len, ceil_div((qo_tile_idx + 1) * CTA_TILE_Q, group_size)); + + DTypeQKAccum s_frag[NUM_MMA_Q][NUM_MMA_KV][8]; + alignas(16) float o_frag[NUM_MMA_Q][NUM_MMA_D_VO][8]; + DTypeQKAccum m[NUM_MMA_Q][2]; + float d[NUM_MMA_Q][2]; + float rope_freq[NUM_MMA_D_QK / 2][4]; + + if constexpr (KTraits::POS_ENCODING_MODE == PosEncodingMode::kRoPELlama) { + const float rope_rcp_scale = params.rope_rcp_scale; + const float rope_rcp_theta = params.rope_rcp_theta; + init_rope_freq(rope_freq, rope_rcp_scale, rope_rcp_theta, tid.x); + } + init_states(variant, o_frag, m, d); + + const uint32_t qo_packed_idx_base = + (qo_tile_idx * NUM_WARPS_Q + get_warp_idx_q(tid.y)) * NUM_MMA_Q * 16; + const uint32_t q_stride_n = params.q_stride_n, q_stride_h = params.q_stride_h; + smem_t qo_smem(smem_storage.q_smem); + const uint32_t o_stride_n = num_qo_heads * HEAD_DIM_VO, o_stride_h = HEAD_DIM_VO; + + DTypeQ* q_ptr_base = + q + q_indptr[request_idx] * q_stride_n + (kv_head_idx * group_size) * q_stride_h; + DTypeO* o_ptr_base = partition_kv ? o + (o_indptr[request_idx] + kv_tile_idx) * o_stride_n + + (kv_head_idx * group_size) * o_stride_h + : o + o_indptr[request_idx] * o_stride_n + + (kv_head_idx * group_size) * o_stride_h; + uint32_t q_smem_offset_r = qo_smem.get_permuted_offset( + get_warp_idx_q(tid.y) * NUM_MMA_Q * 16 + lane_idx % 16, lane_idx / 16); + +#if (__CUDACC_VER_MAJOR__ >= 12 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.wait;"); +#endif + + load_q_global_smem(qo_packed_idx_base, qo_upper_bound, q_ptr_base, q_stride_n, + q_stride_h, group_size, &qo_smem, tid); + + cp_async::commit_group(); + + if constexpr (KTraits::POS_ENCODING_MODE == PosEncodingMode::kRoPELlama) { + cp_async::wait_group<0>(); + block.sync(); + IdType* q_rope_offset = nullptr; + if constexpr (has_maybe_q_rope_offset_v) { + q_rope_offset = params.maybe_q_rope_offset; + } + if (q_rope_offset == nullptr) { + q_smem_inplace_apply_rotary(qo_packed_idx_base, qo_len, kv_len, group_size, + &qo_smem, &q_smem_offset_r, rope_freq, tid); + } else { + q_smem_inplace_apply_rotary_with_pos( + qo_packed_idx_base, q_rope_offset + q_indptr[request_idx], &qo_smem, group_size, + &q_smem_offset_r, rope_freq, tid); + } + block.sync(); + } + + smem_t k_smem(smem_storage.k_smem), v_smem(smem_storage.v_smem); + size_t thr_local_kv_offset[NUM_MMA_KV * KV_THR_LAYOUT_COL / 2 / NUM_WARPS_Q]; + + uint32_t k_smem_offset_r = k_smem.template get_permuted_offset( + get_warp_idx_kv(tid.z) * NUM_MMA_KV * 16 + 8 * (lane_idx / 16) + + lane_idx % 8, + (lane_idx % 16) / 8), + v_smem_offset_r = v_smem.template get_permuted_offset( + get_warp_idx_kv(tid.z) * NUM_MMA_KV * 16 + lane_idx % 16, lane_idx / 16), + k_smem_offset_w = k_smem.template get_permuted_offset( + warp_idx * KV_THR_LAYOUT_ROW + lane_idx / KV_THR_LAYOUT_COL, + lane_idx % KV_THR_LAYOUT_COL), + v_smem_offset_w = v_smem.template get_permuted_offset( + warp_idx * KV_THR_LAYOUT_ROW + lane_idx / KV_THR_LAYOUT_COL, + lane_idx % KV_THR_LAYOUT_COL); + const IdType last_indptr = paged_kv.indptr[paged_kv.batch_size]; + + uint32_t packed_page_iter_base = + paged_kv.indptr[request_idx] * paged_kv.page_size + chunk_start; +#pragma unroll + for (uint32_t i = 0; + i < NUM_MMA_KV * (SWIZZLE_MODE_KV == SwizzleMode::k128B ? 4 : 2) / NUM_WARPS_Q; ++i) { + uint32_t page_iter, entry_idx; + paged_kv.page_size.divmod(packed_page_iter_base + warp_idx * KV_THR_LAYOUT_ROW + + lane_idx / KV_THR_LAYOUT_COL + + KV_THR_LAYOUT_ROW * NUM_WARPS_Q * NUM_WARPS_KV * i, + page_iter, entry_idx); + thr_local_kv_offset[i] = paged_kv.protective_get_kv_offset( + page_iter, kv_head_idx, entry_idx, + (lane_idx % KV_THR_LAYOUT_COL) * upcast_size(), last_indptr); + } + page_produce_kv(&smem_storage, &k_smem_offset_w, paged_kv.k_data, 0, + thr_local_kv_offset, chunk_size, warp_idx, lane_idx); + cp_async::commit_group(); + page_produce_kv(&smem_storage, &v_smem_offset_w, paged_kv.v_data, 0, + thr_local_kv_offset, chunk_size, warp_idx, lane_idx); + cp_async::commit_group(); + + uint32_t num_iterations_prefix; + uint32_t num_iterations_mask; + uint32_t num_iterations = 0; + + if constexpr (MASK_MODE != MaskMode::kMultiItemScoring) { + num_iterations = ceil_div( + (MASK_MODE == MaskMode::kCausal + ? min(chunk_size, + sub_if_greater_or_zero( + kv_len - qo_len + ceil_div(((qo_tile_idx + 1) * CTA_TILE_Q), group_size), + chunk_start)) + : chunk_size), + CTA_TILE_KV); + } else if constexpr (MASK_MODE == MaskMode::kMultiItemScoring) { + num_iterations_prefix = ceil_div( + min(min(chunk_size, + sub_if_greater_or_zero( + kv_len - qo_len + ceil_div(((qo_tile_idx + 1) * CTA_TILE_Q), group_size), + chunk_start)), + sub_if_greater_or_zero(__ldg(maybe_prefix_len_ptr + request_idx), chunk_start)), + CTA_TILE_KV); + num_iterations_mask = + max(min(chunk_size, + sub_if_greater_or_zero( + sub_if_greater_or_zero( + kv_len - qo_len + ceil_div((qo_tile_idx * CTA_TILE_Q), group_size), + __ldg(maybe_max_item_len_ptr + request_idx)), + chunk_start)) / + (CTA_TILE_KV), + num_iterations_prefix); + + num_iterations = max( + num_iterations_mask, + ceil_div(min(chunk_size, + sub_if_greater_or_zero( + kv_len - qo_len + ceil_div(((qo_tile_idx + 1) * CTA_TILE_Q), group_size), + chunk_start)), + CTA_TILE_KV)); + } + + const uint32_t window_iteration = ceil_div( + sub_if_greater_or_zero(kv_len + ceil_div((qo_tile_idx + 1) * CTA_TILE_Q, group_size), + qo_len + window_left + chunk_start), + CTA_TILE_KV); + + const uint32_t mask_iteration = + (MASK_MODE == MaskMode::kCausal || MASK_MODE == MaskMode::kMultiItemScoring + ? min(chunk_size, + sub_if_greater_or_zero( + kv_len + ceil_div((qo_tile_idx * CTA_TILE_Q), group_size) - qo_len, + chunk_start)) + : chunk_size) / + CTA_TILE_KV; + +#pragma unroll 1 + for (uint32_t iter = 0; iter < num_iterations; + iter = (MASK_MODE == MaskMode::kMultiItemScoring) + ? ((iter + 1 == num_iterations_prefix) ? num_iterations_mask : (iter + 1)) + : (iter + 1)) { + const uint32_t prefetch_skip_step = + (MASK_MODE == MaskMode::kMultiItemScoring) + ? ((iter + 1 == num_iterations_prefix) ? (num_iterations_mask - num_iterations_prefix) + : 0) + : 0; + packed_page_iter_base += (1 + prefetch_skip_step) * CTA_TILE_KV; +#pragma unroll + for (uint32_t i = 0; + i < NUM_MMA_KV * (SWIZZLE_MODE_KV == SwizzleMode::k128B ? 4 : 2) / NUM_WARPS_Q; ++i) { + uint32_t page_iter, entry_idx; + paged_kv.page_size.divmod(packed_page_iter_base + warp_idx * KV_THR_LAYOUT_ROW + + lane_idx / KV_THR_LAYOUT_COL + + KV_THR_LAYOUT_ROW * NUM_WARPS_Q * NUM_WARPS_KV * i, + page_iter, entry_idx); + thr_local_kv_offset[i] = paged_kv.protective_get_kv_offset( + page_iter, kv_head_idx, entry_idx, + (lane_idx % KV_THR_LAYOUT_COL) * upcast_size(), last_indptr); + } + cp_async::wait_group<1>(); + block.sync(); + + if constexpr (KTraits::POS_ENCODING_MODE == PosEncodingMode::kRoPELlama) { + k_smem_inplace_apply_rotary( + (paged_kv.rope_pos_offset == nullptr ? 0 : paged_kv.rope_pos_offset[request_idx]) + + chunk_start + iter * CTA_TILE_KV, + &k_smem, &k_smem_offset_r, rope_freq, tid); + block.sync(); + } + + // compute attention score + compute_qk(&qo_smem, &q_smem_offset_r, &k_smem, &k_smem_offset_r, s_frag); + logits_transform( + params, variant, /*batch_idx=*/request_idx, qo_packed_idx_base, + chunk_start + (iter * NUM_WARPS_KV + get_warp_idx_kv(tid.z)) * NUM_MMA_KV * 16, + qo_len, kv_len, group_size, s_frag, tid, kv_head_idx); + + // apply mask + if (MASK_MODE == MaskMode::kCustom) { + logits_mask( + params, variant, /*batch_idx=*/request_idx, qo_packed_idx_base, + chunk_start + (iter * NUM_WARPS_KV + get_warp_idx_kv()) * NUM_MMA_KV * 16, + qo_len, kv_len, chunk_end, group_size, s_frag); + } else { + if constexpr (MASK_MODE != MaskMode::kMultiItemScoring) { + if (iter >= mask_iteration || iter < window_iteration) { + logits_mask( + params, variant, /*batch_idx=*/request_idx, qo_packed_idx_base, + chunk_start + + (iter * NUM_WARPS_KV + get_warp_idx_kv(tid.z)) * NUM_MMA_KV * 16, + qo_len, kv_len, chunk_end, group_size, s_frag); + } + } else if constexpr (MASK_MODE == MaskMode::kMultiItemScoring) { + if (iter + 1 >= num_iterations_prefix) { + logits_mask_multi_item_scoring( + params, variant, /*batch_idx=*/request_idx, qo_packed_idx_base, + chunk_start + + (iter * NUM_WARPS_KV + get_warp_idx_kv(tid.z)) * NUM_MMA_KV * 16, + qo_len, kv_len, window_left, chunk_end, group_size, s_frag, + __ldg(maybe_prefix_len_ptr + request_idx), + maybe_token_pos_in_items_ptr + request_idx * token_pos_in_items_len, tid.x, + kv_head_idx); + } else { + if (iter >= mask_iteration || iter < window_iteration) { + logits_mask( + params, variant, /*batch_idx=*/request_idx, qo_packed_idx_base, + chunk_start + + (iter * NUM_WARPS_KV + get_warp_idx_kv(tid.z)) * NUM_MMA_KV * 16, + qo_len, kv_len, chunk_end, group_size, s_frag); + } + } + } + } + + // compute m,d states in online softmax + update_mdo_states(variant, s_frag, o_frag, m, d); + + block.sync(); + page_produce_kv(&smem_storage, &k_smem_offset_w, paged_kv.k_data, + (iter + 1) * CTA_TILE_KV, thr_local_kv_offset, chunk_size, + warp_idx, lane_idx); + cp_async::commit_group(); + cp_async::wait_group<1>(); + block.sync(); + + // compute sfm*v + compute_sfm_v(&v_smem, &v_smem_offset_r, s_frag, o_frag, d); + + block.sync(); + page_produce_kv(&smem_storage, &v_smem_offset_w, paged_kv.v_data, + (iter + 1) * CTA_TILE_KV, thr_local_kv_offset, chunk_size, + warp_idx, lane_idx); + cp_async::commit_group(); + } + cp_async::wait_group<0>(); + block.sync(); + + finalize_m(variant, m); + + // threadblock synchronization + threadblock_sync_mdo_states(o_frag, &smem_storage, m, d, warp_idx, lane_idx, tid); + + const uint32_t num_kv_chunks = (kv_len_safe + kv_chunk_size - 1) / kv_chunk_size; + + // transform output + transform_output(params, variant, o_frag, m, d, /*batch_idx=*/request_idx, + kv_tile_idx, qo_packed_idx_base, warp_idx, lane_idx, + kv_head_idx, group_size); + + // write_back + write_o_reg_gmem(o_frag, &qo_smem, o_ptr_base, qo_packed_idx_base, qo_len, + /*o_stride_n=*/ + partition_kv ? num_kv_chunks * o_stride_n : o_stride_n, + /*o_stride_h=*/o_stride_h, group_size, tid); + + // write lse + if constexpr (variant.use_softmax) { + if (lse != nullptr) { + if (get_warp_idx_kv(tid.z) == 0) { +#pragma unroll + for (uint32_t mma_q = 0; mma_q < NUM_MMA_Q; ++mma_q) { +#pragma unroll + for (uint32_t j = 0; j < 2; ++j) { + uint32_t q, r; + group_size.divmod(qo_packed_idx_base + lane_idx / 4 + j * 8 + mma_q * 16, q, r); + const uint32_t qo_head_idx = kv_head_idx * group_size + r; + const uint32_t qo_idx = q; + if (qo_idx < qo_upper_bound) { + if (partition_kv) { + lse[(o_indptr[request_idx] + qo_idx * num_kv_chunks + kv_tile_idx) * + num_qo_heads + + qo_head_idx] = math::ptx_log2(d[mma_q][j]) + float(m[mma_q][j]); + } else { + lse[(o_indptr[request_idx] + qo_idx) * num_qo_heads + qo_head_idx] = + math::ptx_log2(d[mma_q][j]) + float(m[mma_q][j]); + } + } + } + } + } + } + } + +#if (__CUDACC_VER_MAJOR__ >= 12 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.launch_dependents;"); +#endif + +#if (__CUDA_ARCH__ < 800) + } +#endif +} + +template +__global__ __launch_bounds__(KTraits::NUM_THREADS) void BatchPrefillWithPagedKVCacheKernel( + const __grid_constant__ Params params) { + extern __shared__ uint8_t smem[]; + auto& smem_storage = reinterpret_cast(smem); + BatchPrefillWithPagedKVCacheDevice(params, smem_storage); +} + +template +cudaError_t BatchPrefillWithRaggedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v, + float* tmp_s, bool enable_pdl, + cudaStream_t stream) { + using DTypeQ = typename Params::DTypeQ; + using DTypeKV = typename Params::DTypeKV; + using DTypeO = typename Params::DTypeO; + const uint32_t padded_batch_size = params.padded_batch_size; + const uint32_t num_qo_heads = params.num_qo_heads; + const uint32_t num_kv_heads = params.num_kv_heads; + constexpr uint32_t NUM_MMA_Q = get_num_mma_q(CTA_TILE_Q); + constexpr uint32_t NUM_WARPS_Q = get_num_warps_q(CTA_TILE_Q); + constexpr uint32_t NUM_WARPS_KV = get_num_warps_kv(CTA_TILE_Q); + + if (padded_batch_size == 0) { + // No request, skip + // this won't happen in CUDAGraph mode because we fixed the padded_batch_size + return cudaSuccess; + } + + dim3 nblks(padded_batch_size, 1, num_kv_heads); + dim3 nthrs(32, NUM_WARPS_Q, NUM_WARPS_KV); + constexpr uint32_t NUM_MMA_D_QK = HEAD_DIM_QK / 16; + constexpr uint32_t NUM_MMA_D_VO = HEAD_DIM_VO / 16; + using DTypeQKAccum = + typename std::conditional, half, + float>::type; + + int dev_id = 0; + FLASHINFER_CUDA_CALL(cudaGetDevice(&dev_id)); + int max_smem_per_sm = 0; + FLASHINFER_CUDA_CALL(cudaDeviceGetAttribute(&max_smem_per_sm, + cudaDevAttrMaxSharedMemoryPerMultiprocessor, dev_id)); + // we expect each sm execute two threadblocks + const int num_ctas_per_sm = + max_smem_per_sm >= 2 * (CTA_TILE_Q * HEAD_DIM_QK * sizeof(DTypeQ) + + (HEAD_DIM_QK + HEAD_DIM_VO) * 16 * NUM_WARPS_KV * sizeof(DTypeKV)) + ? 2 + : 1; + const int max_smem_per_threadblock = max_smem_per_sm / num_ctas_per_sm; + + const uint32_t max_num_mma_kv_reg = + (HEAD_DIM_VO >= 128 && NUM_MMA_Q == 2 && POS_ENCODING_MODE == PosEncodingMode::kRoPELlama && + !USE_FP16_QK_REDUCTION) + ? 2 + : (8 / NUM_MMA_Q); + const uint32_t max_num_mma_kv_smem = + (max_smem_per_threadblock - CTA_TILE_Q * HEAD_DIM_QK * sizeof(DTypeQ)) / + ((HEAD_DIM_QK + HEAD_DIM_VO) * 16 * NUM_WARPS_KV * sizeof(DTypeKV)); + + DISPATCH_NUM_MMA_KV(min(max_num_mma_kv_smem, max_num_mma_kv_reg), NUM_MMA_KV, { + using KTraits = + KernelTraits; + if constexpr (KTraits::IsInvalid()) { + // Invalid configuration, skip + std::ostringstream err_msg; + err_msg << "FlashInfer Internal Error: Invalid configuration : NUM_MMA_Q=" << NUM_MMA_Q + << " NUM_MMA_D_QK=" << NUM_MMA_D_QK << " NUM_MMA_D_VO=" << NUM_MMA_D_VO + << " NUM_MMA_KV=" << NUM_MMA_KV << " NUM_WARPS_Q=" << NUM_WARPS_Q + << " NUM_WARPS_KV=" << NUM_WARPS_KV + << " please create an issue (https://github.com/flashinfer-ai/flashinfer/issues)" + " and report the issue to the developers."; + FLASHINFER_ERROR(err_msg.str()); + } else { + size_t smem_size = sizeof(typename KTraits::SharedStorage); + auto kernel = BatchPrefillWithRaggedKVCacheKernel; + FLASHINFER_CUDA_CALL( + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + // PDL launch config + cudaLaunchAttribute attribute[1]; + cudaLaunchConfig_t config; + if (enable_pdl) { + attribute[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attribute[0].val.programmaticStreamSerializationAllowed = 1; + config.attrs = attribute; + config.numAttrs = 1; + config.gridDim = nblks; + config.blockDim = nthrs; + config.dynamicSmemBytes = smem_size; + config.stream = stream; + } + + if (tmp_v == nullptr) { + // do not partition kv + params.partition_kv = false; + void* args[] = {(void*)¶ms}; + if (enable_pdl) { + FLASHINFER_CUDA_CALL(cudaLaunchKernelEx(&config, kernel, params)); + } else { + FLASHINFER_CUDA_CALL( + cudaLaunchKernel((void*)kernel, nblks, nthrs, args, smem_size, stream)); + } + } else { + // partition kv + params.partition_kv = true; + auto o = params.o; + auto lse = params.lse; + params.o = tmp_v; + params.lse = tmp_s; + void* args[] = {(void*)¶ms}; + if (enable_pdl) { + FLASHINFER_CUDA_CALL(cudaLaunchKernelEx(&config, kernel, params)); + } else { + FLASHINFER_CUDA_CALL( + cudaLaunchKernel((void*)kernel, nblks, nthrs, args, smem_size, stream)); + } + if constexpr (AttentionVariant::use_softmax) { + FLASHINFER_CUDA_CALL(VariableLengthMergeStates( + tmp_v, tmp_s, params.merge_indptr, o, lse, params.max_total_num_rows, + params.total_num_rows, num_qo_heads, HEAD_DIM_VO, enable_pdl, stream)); + } else { + FLASHINFER_CUDA_CALL(VariableLengthAttentionSum( + tmp_v, params.merge_indptr, o, params.max_total_num_rows, params.total_num_rows, + num_qo_heads, HEAD_DIM_VO, enable_pdl, stream)); + } + } + } + }); + return cudaSuccess; +} + +template +cudaError_t BatchPrefillWithPagedKVCacheDispatched(Params params, typename Params::DTypeO* tmp_v, + float* tmp_s, bool enable_pdl, + cudaStream_t stream) { + using DTypeQ = typename Params::DTypeQ; + using DTypeKV = typename Params::DTypeKV; + using DTypeO = typename Params::DTypeO; + const uint32_t padded_batch_size = params.padded_batch_size; + const uint32_t num_qo_heads = params.num_qo_heads; + const uint32_t num_kv_heads = params.paged_kv.num_heads; + constexpr uint32_t NUM_MMA_Q = get_num_mma_q(CTA_TILE_Q); + constexpr uint32_t NUM_WARPS_Q = get_num_warps_q(CTA_TILE_Q); + constexpr uint32_t NUM_WARPS_KV = get_num_warps_kv(CTA_TILE_Q); + + if (padded_batch_size == 0) { + // No request, skip + // this won't happen in CUDAGraph mode because we fixed the padded_batch_size + return cudaSuccess; + } + + dim3 nblks(padded_batch_size, 1, num_kv_heads); + dim3 nthrs(32, NUM_WARPS_Q, NUM_WARPS_KV); + + constexpr uint32_t NUM_MMA_D_QK = HEAD_DIM_QK / 16; + constexpr uint32_t NUM_MMA_D_VO = HEAD_DIM_VO / 16; + using DTypeQKAccum = + typename std::conditional, half, + float>::type; + + int dev_id = 0; + FLASHINFER_CUDA_CALL(cudaGetDevice(&dev_id)); + int max_smem_per_sm = 0; + FLASHINFER_CUDA_CALL(cudaDeviceGetAttribute(&max_smem_per_sm, + cudaDevAttrMaxSharedMemoryPerMultiprocessor, dev_id)); + // we expect each sm execute two threadblocks + const int num_ctas_per_sm = + max_smem_per_sm >= 2 * (CTA_TILE_Q * HEAD_DIM_QK * sizeof(DTypeQ) + + (HEAD_DIM_QK + HEAD_DIM_VO) * 16 * NUM_WARPS_KV * sizeof(DTypeKV)) + ? 2 + : 1; + const int max_smem_per_threadblock = max_smem_per_sm / num_ctas_per_sm; + + const uint32_t max_num_mma_kv_reg = + (HEAD_DIM_VO >= 128 && NUM_MMA_Q == 2 && POS_ENCODING_MODE == PosEncodingMode::kRoPELlama && + !USE_FP16_QK_REDUCTION) + ? 2 + : (8 / NUM_MMA_Q); + const uint32_t max_num_mma_kv_smem = + (max_smem_per_threadblock - CTA_TILE_Q * HEAD_DIM_QK * sizeof(DTypeQ)) / + ((HEAD_DIM_QK + HEAD_DIM_VO) * 16 * NUM_WARPS_KV * sizeof(DTypeKV)); + + DISPATCH_NUM_MMA_KV(min(max_num_mma_kv_smem, max_num_mma_kv_reg), NUM_MMA_KV, { + using KTraits = + KernelTraits; + if constexpr (KTraits::IsInvalid()) { + // Invalid configuration, skip + std::ostringstream err_msg; + err_msg << "FlashInfer Internal Error: Invalid configuration : NUM_MMA_Q=" << NUM_MMA_Q + << " NUM_MMA_D_QK=" << NUM_MMA_D_QK << " NUM_MMA_D_VO=" << NUM_MMA_D_VO + << " NUM_MMA_KV=" << NUM_MMA_KV << " NUM_WARPS_Q=" << NUM_WARPS_Q + << " NUM_WARPS_KV=" << NUM_WARPS_KV + << " please create an issue (https://github.com/flashinfer-ai/flashinfer/issues)" + " and report the issue to the developers."; + FLASHINFER_ERROR(err_msg.str()); + } else { + size_t smem_size = sizeof(typename KTraits::SharedStorage); + auto kernel = BatchPrefillWithPagedKVCacheKernel; + FLASHINFER_CUDA_CALL( + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + // PDL launch config + cudaLaunchAttribute attribute[1]; + cudaLaunchConfig_t config; + if (enable_pdl) { + attribute[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attribute[0].val.programmaticStreamSerializationAllowed = 1; + config.attrs = attribute; + config.numAttrs = 1; + config.gridDim = nblks; + config.blockDim = nthrs; + config.dynamicSmemBytes = smem_size; + config.stream = stream; + } + + if (tmp_v == nullptr) { + // do not partition kv + params.partition_kv = false; + if (enable_pdl) { + FLASHINFER_CUDA_CALL(cudaLaunchKernelEx(&config, kernel, params)); + } else { + void* args[] = {(void*)¶ms}; + FLASHINFER_CUDA_CALL( + cudaLaunchKernel((void*)kernel, nblks, nthrs, args, smem_size, stream)); + } + } else { + params.partition_kv = true; + auto o = params.o; + auto lse = params.lse; + params.o = tmp_v; + params.lse = tmp_s; + if (enable_pdl) { + FLASHINFER_CUDA_CALL(cudaLaunchKernelEx(&config, kernel, params)); + } else { + void* args[] = {(void*)¶ms}; + FLASHINFER_CUDA_CALL( + cudaLaunchKernel((void*)kernel, nblks, nthrs, args, smem_size, stream)); + } + if constexpr (AttentionVariant::use_softmax) { + FLASHINFER_CUDA_CALL(VariableLengthMergeStates( + tmp_v, tmp_s, params.merge_indptr, o, lse, params.max_total_num_rows, + params.total_num_rows, num_qo_heads, HEAD_DIM_VO, enable_pdl, stream)); + } else { + FLASHINFER_CUDA_CALL(VariableLengthAttentionSum( + tmp_v, params.merge_indptr, o, params.max_total_num_rows, params.total_num_rows, + num_qo_heads, HEAD_DIM_VO, enable_pdl, stream)); + } + } + } + }); + return cudaSuccess; +} + +} // namespace flashinfer + +#endif // FLASHINFER_PREFILL_CUH_ diff --git a/include/flashinfer/attention/scheduler.cuh b/include/flashinfer/attention/scheduler.cuh new file mode 100644 index 0000000000000000000000000000000000000000..aaafa2609279ea88ddceab9b80cffe89adb5d709 --- /dev/null +++ b/include/flashinfer/attention/scheduler.cuh @@ -0,0 +1,1691 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef FLASHINFER_ATTENTION_SCHEDULER_CUH_ +#define FLASHINFER_ATTENTION_SCHEDULER_CUH_ + +#include +#include + +#include +#include +#include +#include +#include + +#include "../allocator.h" +#include "../exception.h" +#include "../pos_enc.cuh" +#include "../utils.cuh" +#include "heap.h" + +namespace flashinfer { + +template +__global__ void BatchDecodeWithPagedKVCacheKernel(const __grid_constant__ Params params); + +template +__global__ void BatchDecodeWithPagedKVCacheKernelMLA(Params params); + +template +std::tuple LaunchSpecForDecodeKernelMlaCuteSM80( + const uint32_t num_qo_heads); + +template +__global__ void BatchDecodeWithPagedKVCacheKernelMlaCuteSM80(Params params); + +template +inline void CopyToPageLockedBuffer(void* page_locked_int_buffer, int64_t offset, + const std::vector& vec) { + DType* ptr = GetPtrFromBaseOffset(page_locked_int_buffer, offset); + std::copy(vec.begin(), vec.end(), ptr); +} + +/*! + * \brief Compute the maximum number of pages per batch and the new batch size + * after we partition Paged KV-Cache into multiple chunks on KV sequence length + * dimension. + * \tparam IdType A template type indicates the index data type + * \param max_grid_size The maximum grid size of the kernel + * \param gdy gridDim.y + * \param num_pages The number of pages per request in the batch + * \param max_num_pages_per_batch_lb The pre-set lower bound of maximum number of + * pages per batch, default to 1 + * \return (max_num_pages_per_batch, new_batch_size) The number of pages per batch and + * the new batch size after the partition. + */ +template +inline auto PartitionPagedKVCacheBinarySearchMinNumPagePerBatch( + const uint32_t max_grid_size, const uint32_t gdy, const std::vector& num_pages, + const uint32_t min_num_pages_per_batch = 1) { + uint32_t low = min_num_pages_per_batch, high = 0; + for (const IdType& elem : num_pages) { + high = max(high, elem); + } + uint32_t new_batch_size; + while (low < high) { + uint32_t mid = (low + high) / 2; + new_batch_size = 0; + for (const IdType& elem : num_pages) { + new_batch_size += ceil_div(elem, mid); + } + if (new_batch_size * gdy > max_grid_size) { + low = mid + 1; + } else { + high = mid; + } + } + new_batch_size = 0; + for (const IdType& elem : num_pages) { + new_batch_size += ceil_div(std::max(elem, 1), low); + } + return std::make_tuple(low, new_batch_size); +} + +inline auto PrefillBinarySearchKVChunkSize(const bool enable_cuda_graph, + const uint32_t max_batch_size_if_split, + const std::vector& packed_qo_len_arr, + const std::vector& kv_len_arr, + const uint32_t qo_chunk_size, + const uint32_t min_kv_chunk_size = 1) { + const int64_t batch_size = packed_qo_len_arr.size(); + int64_t max_kv_len = 1; + for (const int64_t& kv_len : kv_len_arr) { + max_kv_len = std::max(max_kv_len, kv_len); + } + + int64_t low = min_kv_chunk_size; + int64_t high = max_kv_len; + constexpr int64_t min_kv_len = 1; + while (low < high) { + const int64_t mid = (low + high) / 2; + int64_t new_batch_size = 0; + for (uint32_t i = 0; i < batch_size; ++i) { + new_batch_size += ceil_div(packed_qo_len_arr[i], qo_chunk_size) * + ceil_div(std::max(kv_len_arr[i], min_kv_len), mid); + } + if (new_batch_size > max_batch_size_if_split) { + low = mid + 1; + } else { + high = mid; + } + } + return std::make_tuple(enable_cuda_graph || low < max_kv_len, low); +} + +/*! + * \brief Estimate the temporary buffer size and the maximum grid size for the + * partition-kv BatchDecodeWithPagedKVCache kernel + * \tparam DTypeKV A template type indicates the key-value data type + * \tparam DTypeO A template type indicates the output data type + * \tparam IdType A template type indicates the index data type + * \param split_kv Whether to split the KV cache into multiple chunks + * \param max_grid_size The maximum grid size that can be used in a partiton-kv kernel + * \param max_num_pages_per_batch The maximum number of pages per batch + * \param new_batch_size The new batch size after the partition + * \param paged_kv The paged kv cache data structure + * \param num_qo_heads A integer indicates the number of heads of query and output + * \param pos_encoding_mode The positional encoding mode + * \param stream The cuda stream to launch the kernel + * \return status Indicates whether CUDA calls are successful + */ +template +inline cudaError_t BatchDecodeWithPagedKVCacheWorkEstimationDispatched( + bool& split_kv, uint32_t& max_grid_size, uint32_t& max_num_pages_per_batch, + uint32_t& new_batch_size, uint32_t& gdy, uint32_t batch_size, + typename Params::IdType* kv_indptr_h, const uint32_t num_qo_heads, const uint32_t page_size, + bool enable_cuda_graph, cudaStream_t stream) { + using DTypeKV = typename Params::DTypeKV; + using IdType = typename Params::IdType; + constexpr uint32_t vec_size = std::max(16UL / sizeof(DTypeKV), HEAD_DIM / 32UL); + auto compute_capacity = GetCudaComputeCapability(); + DISPATCH_COMPUTE_CAP_DECODE_NUM_STAGES_SMEM(compute_capacity, NUM_STAGES_SMEM, { + constexpr uint32_t bdx = HEAD_DIM / vec_size; + static_assert(bdx <= 32); + constexpr uint32_t bdy = GROUP_SIZE; + constexpr uint32_t num_threads = std::max(128U, bdx * bdy); + constexpr uint32_t bdz = num_threads / (bdx * bdy); + constexpr uint32_t tile_size_per_bdx = GROUP_SIZE == 1 ? (sizeof(DTypeKV) == 1 ? 2U : 4U) : 1U; + const uint32_t num_kv_heads = num_qo_heads / GROUP_SIZE; + gdy = num_kv_heads; + const uint32_t smem_size = + 2 * NUM_STAGES_SMEM * tile_size_per_bdx * bdy * bdz * HEAD_DIM * sizeof(DTypeKV) + + std::max(tile_size_per_bdx * num_threads * sizeof(DTypeKV*), 2 * bdy * bdz * sizeof(float)); + + auto kernel = + BatchDecodeWithPagedKVCacheKernel; + int num_blocks_per_sm = 0; + int num_sm = 0; + int dev_id = 0; + FLASHINFER_CUDA_CALL(cudaGetDevice(&dev_id)); + FLASHINFER_CUDA_CALL(cudaDeviceGetAttribute(&num_sm, cudaDevAttrMultiProcessorCount, dev_id)); + FLASHINFER_CUDA_CALL(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&num_blocks_per_sm, kernel, + num_threads, smem_size)); + max_grid_size = num_blocks_per_sm * num_sm; + if (batch_size * gdy >= max_grid_size) { + split_kv = false; + max_num_pages_per_batch = 1; + for (uint32_t batch_idx = 0; batch_idx < batch_size; ++batch_idx) { + max_num_pages_per_batch = std::max( + max_num_pages_per_batch, kv_indptr_h[batch_idx + 1] - kv_indptr_h[batch_idx]); + } + new_batch_size = batch_size; + } else { + // compute max_num_pages_per_batch and new_batch_size + std::vector num_pages(batch_size); + for (uint32_t batch_idx = 0; batch_idx < batch_size; ++batch_idx) { + num_pages[batch_idx] = kv_indptr_h[batch_idx + 1] - kv_indptr_h[batch_idx]; + } + std::tie(max_num_pages_per_batch, new_batch_size) = + PartitionPagedKVCacheBinarySearchMinNumPagePerBatch(max_grid_size, gdy, num_pages, + std::max(128 / page_size, 1U)); + if (new_batch_size == batch_size && !enable_cuda_graph) { + // do not use partition-kv kernel for short sequence, when not using CUDAGraph + split_kv = false; + } else { + // when using CUDAGraph, we always use partition-kv kernel + split_kv = true; + } + } + return cudaSuccess; + }) +} + +template +inline cudaError_t BatchDecodeWithPagedKVCacheWorkEstimationDispatchedMLA( + bool& split_kv, uint32_t& max_grid_size, uint32_t& max_num_pages_per_batch, + uint32_t& new_batch_size, uint32_t& gdy, uint32_t batch_size, + typename Params::IdType* kv_indptr_h, const uint32_t num_qo_heads, const uint32_t page_size, + bool enable_cuda_graph, cudaStream_t stream) { + using DTypeKV = typename Params::DTypeKV; + using IdType = typename Params::IdType; + + auto compute_capacity = GetCudaComputeCapability(); + DISPATCH_COMPUTE_CAP_DECODE_NUM_STAGES_SMEM(compute_capacity, NUM_STAGES_SMEM, { + constexpr uint32_t vec_size_ckv = std::max(16UL / sizeof(DTypeKV), HEAD_DIM_CKV / 32UL); + constexpr uint32_t bdx = HEAD_DIM_CKV / vec_size_ckv; + constexpr uint32_t vec_size_kpe = HEAD_DIM_KPE / bdx; + + constexpr uint32_t bdy = 8; + constexpr uint32_t tile_size_qo_heads = 2; + constexpr uint32_t qo_heads_per_block = bdy * tile_size_qo_heads; + constexpr uint32_t num_threads = std::max(128U, bdx * bdy); + constexpr uint32_t bdz = num_threads / (bdx * bdy); + gdy = ceil_div(num_qo_heads, qo_heads_per_block); + + const uint32_t smem_size = + NUM_STAGES_SMEM * bdy * bdz * (HEAD_DIM_CKV + HEAD_DIM_KPE) * sizeof(DTypeKV) + + std::max(num_threads * sizeof(size_t) * 2, 2 * bdy * bdz * sizeof(float)); + + auto kernel = + BatchDecodeWithPagedKVCacheKernelMLA; + int num_blocks_per_sm = 0; + int num_sm = 0; + int dev_id = 0; + FLASHINFER_CUDA_CALL(cudaGetDevice(&dev_id)); + FLASHINFER_CUDA_CALL(cudaDeviceGetAttribute(&num_sm, cudaDevAttrMultiProcessorCount, dev_id)); + FLASHINFER_CUDA_CALL(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&num_blocks_per_sm, kernel, + num_threads, smem_size)); + max_grid_size = num_blocks_per_sm * num_sm; + if (batch_size * gdy >= max_grid_size) { + split_kv = false; + max_num_pages_per_batch = 1; + for (uint32_t batch_idx = 0; batch_idx < batch_size; ++batch_idx) { + max_num_pages_per_batch = std::max( + max_num_pages_per_batch, kv_indptr_h[batch_idx + 1] - kv_indptr_h[batch_idx]); + } + new_batch_size = batch_size; + } else { + // compute max_num_pages_per_batch and new_batch_size + std::vector num_pages(batch_size); + for (uint32_t batch_idx = 0; batch_idx < batch_size; ++batch_idx) { + num_pages[batch_idx] = kv_indptr_h[batch_idx + 1] - kv_indptr_h[batch_idx]; + } + std::tie(max_num_pages_per_batch, new_batch_size) = + PartitionPagedKVCacheBinarySearchMinNumPagePerBatch(max_grid_size, gdy, num_pages, + std::max(128 / page_size, 1U)); + if (new_batch_size == batch_size && !enable_cuda_graph) { + // do not use partition-kv kernel for short sequence, when not using CUDAGraph + split_kv = false; + } else { + // when using CUDAGraph, we always use partition-kv kernel + split_kv = true; + } + } + + return cudaSuccess; + }); +} + +template +inline cudaError_t BatchDecodeWithPagedKVCacheWorkEstimationDispatchedMlaCuteSM80( + bool& split_kv, uint32_t& max_grid_size, uint32_t& max_num_pages_per_batch, + uint32_t& new_batch_size, uint32_t& gdy_, uint32_t batch_size, + typename Params::IdType* kv_indptr_h, const uint32_t num_qo_heads, const uint32_t page_size, + bool enable_cuda_graph, cudaStream_t stream) { + using DTypeKV = typename Params::DTypeKV; + using IdType = typename Params::IdType; + + auto [smem_size, gdy, k_warps] = + LaunchSpecForDecodeKernelMlaCuteSM80( + num_qo_heads); + gdy_ = gdy; + const uint32_t num_threads = k_warps * 32; + auto kernel = + BatchDecodeWithPagedKVCacheKernelMlaCuteSM80; + int num_blocks_per_sm; + int num_sm = 0; + int dev_id = 0; + FLASHINFER_CUDA_CALL(cudaGetDevice(&dev_id)); + FLASHINFER_CUDA_CALL(cudaDeviceGetAttribute(&num_sm, cudaDevAttrMultiProcessorCount, dev_id)); + + // FLASHINFER_CUDA_CALL(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&num_blocks_per_sm, kernel, + // num_threads, smem_size)); + // fixme: num_blocks_per_sm is 0 derived from cudaOccupancyMaxActiveBlocksPerMultiprocessor at + // times, and we fill smem with q-heads as many as possible, so num_blocks_per_sm should be 1 + num_blocks_per_sm = 1; + + max_grid_size = num_blocks_per_sm * num_sm; + if (batch_size * gdy >= max_grid_size) { + split_kv = false; + max_num_pages_per_batch = 1; + for (uint32_t batch_idx = 0; batch_idx < batch_size; ++batch_idx) { + max_num_pages_per_batch = std::max( + max_num_pages_per_batch, kv_indptr_h[batch_idx + 1] - kv_indptr_h[batch_idx]); + } + new_batch_size = batch_size; + } else { + // compute max_num_pages_per_batch and new_batch_size + std::vector num_pages(batch_size); + for (uint32_t batch_idx = 0; batch_idx < batch_size; ++batch_idx) { + num_pages[batch_idx] = kv_indptr_h[batch_idx + 1] - kv_indptr_h[batch_idx]; + } + std::tie(max_num_pages_per_batch, new_batch_size) = + PartitionPagedKVCacheBinarySearchMinNumPagePerBatch(max_grid_size, gdy, num_pages, + std::max(128 / page_size, 1U)); + if (new_batch_size == batch_size && !enable_cuda_graph) { + // do not use partition-kv kernel for short sequence, when not using CUDAGraph + split_kv = false; + } else { + // when using CUDAGraph, we always use partition-kv kernel + split_kv = true; + } + } + + return cudaSuccess; +} + +/*! + * \brief Partition Paged KV-Cache into multiple chunks on KV sequence length + * \tparam IdType A template type indicates the index data type + * \param old_batch_size The batch size of the old Paged KV-Cache + * \param old_page_indptr_h The host-side page indptr of the old Paged KV-Cache + * \param max_num_pages_per_batch The maximum number of pages per batch + * \param new_paged_kv_d The device-side new Paged KV-Cache + * \param stream The cuda stream to launch the kernel + * \return status Indicates whether CUDA calls are successful + */ +template +inline auto DecodeSplitKVIndptr(IdType* indptr_h, uint32_t batch_size, uint32_t kv_chunk_size) { + std::vector request_indices, kv_tile_indices, o_indptr; + o_indptr.push_back(0); + + for (uint32_t batch_idx = 0; batch_idx < batch_size; batch_idx++) { + uint32_t num_tiles_kv = ceil_div( + std::max(indptr_h[batch_idx + 1] - indptr_h[batch_idx], 1U), kv_chunk_size); + for (uint32_t kv_tile_idx = 0; kv_tile_idx < num_tiles_kv; ++kv_tile_idx) { + request_indices.push_back(batch_idx); + kv_tile_indices.push_back(kv_tile_idx); + } + o_indptr.push_back(o_indptr.back() + num_tiles_kv); + } + + return std::make_tuple(request_indices, kv_tile_indices, o_indptr); +} + +struct DecodePlanInfo { + int64_t padded_batch_size; + int64_t v_offset; + int64_t s_offset; + int64_t request_indices_offset; + int64_t kv_tile_indices_offset; + int64_t o_indptr_offset; + int64_t block_valid_mask_offset; + int64_t kv_chunk_size_ptr_offset; + bool enable_cuda_graph; + bool split_kv; + + DecodePlanInfo() + : padded_batch_size(0), + v_offset(0), + s_offset(0), + request_indices_offset(0), + kv_tile_indices_offset(0), + o_indptr_offset(0), + block_valid_mask_offset(0), + kv_chunk_size_ptr_offset(0), + enable_cuda_graph(false), + split_kv(false) {} + + // convert DecodePlanInfo to std::vector + std::vector ToVector() const { + return {padded_batch_size, + v_offset, + s_offset, + request_indices_offset, + kv_tile_indices_offset, + o_indptr_offset, + block_valid_mask_offset, + kv_chunk_size_ptr_offset, + enable_cuda_graph, + split_kv}; + } + + // From std::vector to DecodePlanInfo + void FromVector(const std::vector& vec) { + if (vec.size() != 10) { + std::ostringstream err_msg; + err_msg << "DecodePlanInfo::FromVector: vec.size() should be 10, but got " << vec.size(); + FLASHINFER_ERROR(err_msg.str()); + } + padded_batch_size = vec[0]; + v_offset = vec[1]; + s_offset = vec[2]; + request_indices_offset = vec[3]; + kv_tile_indices_offset = vec[4]; + o_indptr_offset = vec[5]; + block_valid_mask_offset = vec[6]; + kv_chunk_size_ptr_offset = vec[7]; + enable_cuda_graph = vec[8]; + split_kv = vec[9]; + } +}; + +template +inline cudaError_t DecodePlan(void* float_buffer, size_t float_workspace_size_in_bytes, + void* int_buffer, void* page_locked_int_buffer, + size_t int_workspace_size_in_bytes, DecodePlanInfo& plan_info, + typename Params::IdType* indptr_h, uint32_t batch_size, + uint32_t num_qo_heads, uint32_t page_size, bool enable_cuda_graph, + cudaStream_t stream, WorkEstimationFunc work_estimation_func) { + using DTypeO = typename Params::DTypeO; + using IdType = typename Params::IdType; + bool split_kv; + uint32_t max_grid_size, kv_chunk_size_in_pages, new_batch_size, gdy; + + FLASHINFER_CUDA_CALL(work_estimation_func(split_kv, max_grid_size, kv_chunk_size_in_pages, + new_batch_size, gdy, batch_size, indptr_h, num_qo_heads, + page_size, enable_cuda_graph, stream)); + size_t padded_batch_size; + plan_info.enable_cuda_graph = enable_cuda_graph; + plan_info.split_kv = split_kv; + padded_batch_size = + (enable_cuda_graph) ? (split_kv ? max_grid_size / gdy : batch_size) : new_batch_size; + plan_info.padded_batch_size = padded_batch_size; + + auto [request_indices_vec, kv_tile_indices_vec, o_indptr_vec] = + DecodeSplitKVIndptr(indptr_h, batch_size, kv_chunk_size_in_pages); + + AlignedAllocator int_allocator(int_buffer, int_workspace_size_in_bytes); + plan_info.request_indices_offset = int_allocator.aligned_alloc_offset( + padded_batch_size * sizeof(IdType), 16, "batch_decode_request_indices"); + plan_info.kv_tile_indices_offset = int_allocator.aligned_alloc_offset( + padded_batch_size * sizeof(IdType), 16, "batch_decode_kv_tile_indices"); + plan_info.o_indptr_offset = int_allocator.aligned_alloc_offset( + (padded_batch_size + 1) * sizeof(IdType), 16, "batch_decode_o_indptr"); + plan_info.kv_chunk_size_ptr_offset = + int_allocator.aligned_alloc_offset(sizeof(IdType), 1, "batch_decode_kv_chunk_size_ptr"); + IdType* request_indices_h = + GetPtrFromBaseOffset(page_locked_int_buffer, plan_info.request_indices_offset); + IdType* kv_tile_indices_h = + GetPtrFromBaseOffset(page_locked_int_buffer, plan_info.kv_tile_indices_offset); + IdType* o_indptr_h = + GetPtrFromBaseOffset(page_locked_int_buffer, plan_info.o_indptr_offset); + IdType* kv_chunk_size_ptr_h = + GetPtrFromBaseOffset(page_locked_int_buffer, plan_info.kv_chunk_size_ptr_offset); + std::copy(request_indices_vec.begin(), request_indices_vec.end(), request_indices_h); + std::copy(kv_tile_indices_vec.begin(), kv_tile_indices_vec.end(), kv_tile_indices_h); + std::copy(o_indptr_vec.begin(), o_indptr_vec.end(), o_indptr_h); + kv_chunk_size_ptr_h[0] = kv_chunk_size_in_pages * page_size; + + if (split_kv) { + AlignedAllocator float_allocator(float_buffer, float_workspace_size_in_bytes); + plan_info.v_offset = float_allocator.aligned_alloc_offset( + num_qo_heads * padded_batch_size * HEAD_DIM * sizeof(float), 16, "batch_decode_tmp_v"); + plan_info.s_offset = float_allocator.aligned_alloc_offset( + num_qo_heads * padded_batch_size * sizeof(float), 16, "batch_decode_tmp_s"); + + plan_info.block_valid_mask_offset = int_allocator.aligned_alloc_offset( + padded_batch_size * sizeof(bool), 16, "batch_decode_block_valid_mask"); + bool* block_valid_mask_h = + GetPtrFromBaseOffset(page_locked_int_buffer, plan_info.block_valid_mask_offset); + for (uint32_t i = 0; i < padded_batch_size; ++i) { + block_valid_mask_h[i] = i < new_batch_size; + } + } + + size_t num_bytes_to_copy = int_allocator.num_allocated_bytes(); + + FLASHINFER_CUDA_CALL(cudaMemcpyAsync(int_buffer, page_locked_int_buffer, num_bytes_to_copy, + cudaMemcpyHostToDevice, stream)); + return cudaSuccess; +} + +template +inline auto PrefillSplitQOKVIndptr(IdType* qo_indptr_h, IdType* kv_indptr_h, + uint32_t total_num_rows, uint32_t batch_size, + uint32_t num_qo_heads, uint32_t num_kv_heads, uint32_t head_dim, + uint32_t page_size, uint32_t max_batch_size_if_split, + bool enable_cuda_graph) { + std::vector request_indices, qo_tile_indices, kv_tile_indices, merge_indptr, o_indptr; + merge_indptr.push_back(0); + o_indptr.push_back(0); + + const uint32_t gqa_group_size = num_qo_heads / num_kv_heads; + + // step 1: determine packed_qo_len_arr and verify qo_indptr contents. + std::vector packed_qo_len_arr(batch_size), kv_len_arr(batch_size); + for (uint32_t i = 0; i < batch_size; ++i) { + packed_qo_len_arr[i] = int64_t(qo_indptr_h[i + 1] - qo_indptr_h[i]) * int64_t(gqa_group_size); + if (packed_qo_len_arr[i] < 0) { + std::ostringstream err_msg; + err_msg << "qo_indptr[" << i + 1 << "]" << qo_indptr_h[i + 1] << " - qo_indptr[" << i << "]" + << qo_indptr_h[i] << " should be non-negative"; + FLASHINFER_ERROR(err_msg.str()); + } + kv_len_arr[i] = int64_t(kv_indptr_h[i + 1] - kv_indptr_h[i]); + if (kv_len_arr[i] < 0) { + std::ostringstream err_msg; + err_msg << "kv_indptr[" << i + 1 << "]" << kv_indptr_h[i + 1] << " - kv_indptr[" << i << "]" + << kv_indptr_h[i] << " should be non-negative"; + FLASHINFER_ERROR(err_msg.str()); + } + } + + // step 2: determine cta_tile_q, kv_chunk_size and total_num_tiles_q + const uint32_t min_kv_chunk_size = std::max((128 / page_size), 1U); + uint32_t cta_tile_q; + uint32_t total_num_tiles_q; + if (enable_cuda_graph) { + // When CUDA graphs are enabled, the lengths of sequences determined by + // qo_indptr_h can vary. We assume that the dummy data based on which + // the CUDA graph is created fixes the maximum number of tokens. + const uint64_t max_seq_len = total_num_rows - batch_size + 1; + uint64_t max_qo_len = uint64_t(max_seq_len) * gqa_group_size; + cta_tile_q = FA2DetermineCtaTileQ(max_qo_len, head_dim); + + // Find an upper bound for the number of tiles, derived from the total + // number of rows and the batch size. The sum of qo lengths rounded + // up to cta_tile_q will not exceed this number derived from the total + // number of rows. + total_num_tiles_q = ceil_div(total_num_rows * gqa_group_size, cta_tile_q) + batch_size - 1; + } else { + int64_t sum_packed_qo_len = 0; + for (uint32_t i = 0; i < batch_size; ++i) { + sum_packed_qo_len += packed_qo_len_arr[i]; + } + const int64_t avg_packed_qo_len = sum_packed_qo_len / batch_size; + cta_tile_q = FA2DetermineCtaTileQ(avg_packed_qo_len, head_dim); + + total_num_tiles_q = 0; + for (uint32_t i = 0; i < batch_size; ++i) { + total_num_tiles_q += ceil_div(packed_qo_len_arr[i], cta_tile_q); + } + } + + auto [split_kv, kv_chunk_size] = + PrefillBinarySearchKVChunkSize(enable_cuda_graph, max_batch_size_if_split, packed_qo_len_arr, + kv_len_arr, cta_tile_q, min_kv_chunk_size); + + // step 3: split qo_indptr and kv_indptr + uint32_t new_batch_size = 0; + for (uint32_t request_idx = 0; request_idx < batch_size; ++request_idx) { + const int64_t packed_qo_len = packed_qo_len_arr[request_idx]; + const int64_t kv_len = std::max(int(kv_len_arr[request_idx]), 1); + const int64_t num_tiles_q = ceil_div(packed_qo_len, cta_tile_q); + const int64_t num_tiles_kv = ceil_div(kv_len, kv_chunk_size); + + for (uint32_t q_tile_idx = 0; q_tile_idx < num_tiles_q; ++q_tile_idx) { + for (uint32_t kv_tile_idx = 0; kv_tile_idx < num_tiles_kv; ++kv_tile_idx) { + new_batch_size += 1; + request_indices.push_back(request_idx); + qo_tile_indices.push_back(q_tile_idx); + kv_tile_indices.push_back(kv_tile_idx); + } + } + + int64_t qo_len = packed_qo_len / gqa_group_size; + for (uint32_t row = 0; row < qo_len; ++row) { + merge_indptr.push_back(merge_indptr.back() + num_tiles_kv); + } + o_indptr.push_back(o_indptr.back() + qo_len * num_tiles_kv); + } + + const size_t padded_batch_size = + enable_cuda_graph ? std::max(max_batch_size_if_split, total_num_tiles_q) : new_batch_size; + FLASHINFER_CHECK(new_batch_size <= padded_batch_size, + "new batch size should not exceed padded batch size"); + + // step 4: multiply kv_chunk_size by page_size + kv_chunk_size *= page_size; + + return std::make_tuple(split_kv, new_batch_size, padded_batch_size, cta_tile_q, kv_chunk_size, + std::move(request_indices), std::move(qo_tile_indices), + std::move(kv_tile_indices), std::move(merge_indptr), std::move(o_indptr)); +} + +struct PrefillPlanInfo { + int64_t padded_batch_size; + int64_t total_num_rows; + int64_t total_num_rows_offset; + int64_t cta_tile_q; + int64_t request_indices_offset; + int64_t qo_tile_indices_offset; + int64_t kv_tile_indices_offset; + int64_t merge_indptr_offset; + int64_t o_indptr_offset; + int64_t kv_chunk_size_ptr_offset; + int64_t v_offset; + int64_t s_offset; + int64_t block_valid_mask_offset; + bool enable_cuda_graph; + bool split_kv; + + PrefillPlanInfo() + : padded_batch_size(0), + total_num_rows(0), + total_num_rows_offset(0), + cta_tile_q(0), + request_indices_offset(0), + qo_tile_indices_offset(0), + kv_tile_indices_offset(0), + merge_indptr_offset(0), + o_indptr_offset(0), + kv_chunk_size_ptr_offset(0), + v_offset(0), + s_offset(0), + block_valid_mask_offset(0), + enable_cuda_graph(false), + split_kv(false) {} + + // convert PrefillPlanInfo to std::vector + std::vector ToVector() const { + return {padded_batch_size, + total_num_rows, + total_num_rows_offset, + cta_tile_q, + request_indices_offset, + qo_tile_indices_offset, + kv_tile_indices_offset, + merge_indptr_offset, + o_indptr_offset, + kv_chunk_size_ptr_offset, + v_offset, + s_offset, + block_valid_mask_offset, + enable_cuda_graph, + split_kv}; + } + + // From std::vector to PrefillPlanInfo + void FromVector(const std::vector& vec) { + if (vec.size() != 15) { + std::ostringstream err_msg; + err_msg << "PrefillPlanInfo::FromVector: vec.size() should be 15, but got " << vec.size(); + FLASHINFER_ERROR(err_msg.str()); + } + padded_batch_size = vec[0]; + total_num_rows = vec[1]; + total_num_rows_offset = vec[2]; + cta_tile_q = vec[3]; + request_indices_offset = vec[4]; + qo_tile_indices_offset = vec[5]; + kv_tile_indices_offset = vec[6]; + merge_indptr_offset = vec[7]; + o_indptr_offset = vec[8]; + kv_chunk_size_ptr_offset = vec[9]; + v_offset = vec[10]; + s_offset = vec[11]; + block_valid_mask_offset = vec[12]; + enable_cuda_graph = vec[13]; + split_kv = vec[14]; + } +}; + +template +inline cudaError_t PrefillPlan(void* float_buffer, size_t float_workspace_size_in_bytes, + void* int_buffer, void* page_locked_int_buffer, + size_t int_workspace_size_in_bytes, PrefillPlanInfo& plan_info, + IdType* qo_indptr_h, IdType* kv_indptr_h, uint32_t total_num_rows, + uint32_t batch_size, uint32_t num_qo_heads, uint32_t num_kv_heads, + uint32_t head_dim_qk, uint32_t head_dim_vo, uint32_t page_size, + bool enable_cuda_graph, uint32_t sizeof_dtype_o, + cudaStream_t stream) { + if (num_qo_heads % num_kv_heads != 0) { + std::ostringstream err_msg; + err_msg << "num_qo_heads " << num_qo_heads << " should be divisible by num_kv_heads " + << num_kv_heads; + FLASHINFER_ERROR(err_msg.str()); + } + + // step 0: get the number of SMs + int num_sm = 0; + int dev_id = 0; + FLASHINFER_CUDA_CALL(cudaGetDevice(&dev_id)); + FLASHINFER_CUDA_CALL(cudaDeviceGetAttribute(&num_sm, cudaDevAttrMultiProcessorCount, dev_id)); + int num_blocks_per_sm = 2; + int max_grid_size = num_blocks_per_sm * num_sm; + uint32_t max_batch_size_if_split = max_grid_size / num_kv_heads; + + // step 2: determine kv_chunk_size + auto [split_kv, new_batch_size, padded_batch_size, cta_tile_q, kv_chunk_size, request_indices_vec, + qo_tile_indices_vec, kv_tile_indices_vec, merge_indptr_vec, o_indptr_vec] = + PrefillSplitQOKVIndptr(qo_indptr_h, kv_indptr_h, total_num_rows, batch_size, num_qo_heads, + num_kv_heads, head_dim_vo, page_size, max_batch_size_if_split, + enable_cuda_graph); + + plan_info.cta_tile_q = cta_tile_q; + plan_info.total_num_rows = total_num_rows; + plan_info.enable_cuda_graph = enable_cuda_graph; + plan_info.padded_batch_size = padded_batch_size; + plan_info.split_kv = split_kv; + + AlignedAllocator int_allocator(int_buffer, int_workspace_size_in_bytes); + plan_info.request_indices_offset = int_allocator.aligned_alloc_offset( + sizeof(IdType) * padded_batch_size, 16, "batch_prefill_request_indices"); + plan_info.qo_tile_indices_offset = int_allocator.aligned_alloc_offset( + sizeof(IdType) * padded_batch_size, 16, "batch_prefill_qo_tile_indices"); + plan_info.kv_tile_indices_offset = int_allocator.aligned_alloc_offset( + sizeof(IdType) * padded_batch_size, 16, "batch_prefill_kv_tile_indices"); + plan_info.o_indptr_offset = int_allocator.aligned_alloc_offset(sizeof(IdType) * (batch_size + 1), + 16, "batch_prefill_o_indptr"); + plan_info.kv_chunk_size_ptr_offset = + int_allocator.aligned_alloc_offset(sizeof(IdType), 1, "batch_prefill_kv_chunk_size_ptr"); + + if (plan_info.enable_cuda_graph) { + plan_info.total_num_rows_offset = + int_allocator.aligned_alloc_offset(sizeof(uint32_t), 16, "batch_prefill_total_num_rows"); + uint32_t* total_num_rows_h = + GetPtrFromBaseOffset(page_locked_int_buffer, plan_info.total_num_rows_offset); + *total_num_rows_h = qo_indptr_h[batch_size]; + } + + IdType* request_indices_h = + GetPtrFromBaseOffset(page_locked_int_buffer, plan_info.request_indices_offset); + IdType* qo_tile_indices_h = + GetPtrFromBaseOffset(page_locked_int_buffer, plan_info.qo_tile_indices_offset); + IdType* kv_tile_indices_h = + GetPtrFromBaseOffset(page_locked_int_buffer, plan_info.kv_tile_indices_offset); + IdType* o_indptr_h = + GetPtrFromBaseOffset(page_locked_int_buffer, plan_info.o_indptr_offset); + IdType* kv_chunk_size_ptr_h = + GetPtrFromBaseOffset(page_locked_int_buffer, plan_info.kv_chunk_size_ptr_offset); + std::copy(request_indices_vec.begin(), request_indices_vec.end(), request_indices_h); + std::copy(qo_tile_indices_vec.begin(), qo_tile_indices_vec.end(), qo_tile_indices_h); + std::copy(kv_tile_indices_vec.begin(), kv_tile_indices_vec.end(), kv_tile_indices_h); + std::copy(o_indptr_vec.begin(), o_indptr_vec.end(), o_indptr_h); + kv_chunk_size_ptr_h[0] = kv_chunk_size; + + if (split_kv) { + AlignedAllocator float_allocator(float_buffer, float_workspace_size_in_bytes); + plan_info.v_offset = float_allocator.aligned_alloc_offset( + num_qo_heads * padded_batch_size * cta_tile_q * head_dim_vo * sizeof(float), 16, + "batch_prefill_tmp_v"); + plan_info.s_offset = float_allocator.aligned_alloc_offset( + num_qo_heads * padded_batch_size * cta_tile_q * sizeof(float), 16, "batch_prefill_tmp_s"); + plan_info.merge_indptr_offset = int_allocator.aligned_alloc_offset( + sizeof(IdType) * (plan_info.total_num_rows + 1), 16, "batch_prefill_merge_indptr"); + plan_info.block_valid_mask_offset = int_allocator.aligned_alloc_offset( + sizeof(bool) * padded_batch_size, 16, "batch_prefill_block_valid_mask"); + + IdType* merge_indptr_h = + GetPtrFromBaseOffset(page_locked_int_buffer, plan_info.merge_indptr_offset); + bool* block_valid_mask_h = + GetPtrFromBaseOffset(page_locked_int_buffer, plan_info.block_valid_mask_offset); + std::copy(merge_indptr_vec.begin(), merge_indptr_vec.end(), merge_indptr_h); + for (uint32_t i = 0; i < padded_batch_size; ++i) { + block_valid_mask_h[i] = i < new_batch_size; + } + } + + size_t num_bytes_to_copy = int_allocator.num_allocated_bytes(); + FLASHINFER_CUDA_CALL(cudaMemcpyAsync(int_buffer, page_locked_int_buffer, num_bytes_to_copy, + cudaMemcpyHostToDevice, stream)); + + return cudaSuccess; +} + +inline float cost_function(int qo_len, int kv_len) { return 2 * float(qo_len) + kv_len; } + +template +std::vector flatten(const std::vector>& vec, int size_after_flatten) { + std::vector result; + result.reserve(size_after_flatten); + for (const auto& inner_vec : vec) { + result.insert(result.end(), inner_vec.begin(), inner_vec.end()); + } + return result; +} + +inline int packed_causal_kv_end(int qo_len, int kv_len, int qo_tile_idx, int cluster_tile_q, + int num_qo_tiles, int group_size) { + if (qo_tile_idx + 1 == num_qo_tiles) { + return kv_len; + } + int kv_len_init = kv_len - qo_len; // right aligned + return min(kv_len_init + ceil_div((qo_tile_idx + 1) * cluster_tile_q, group_size), kv_len); +} + +struct PrefillPlanSM90Info { + int64_t qo_tile_indices_offset; + int64_t qo_indptr_offset; + int64_t kv_indptr_offset; + int64_t qo_len_offset; + int64_t kv_len_offset; + int64_t head_indices_offset; + int64_t work_indptr_offset; + int64_t batch_indices_offset; + bool same_schedule_for_all_heads; + + PrefillPlanSM90Info() + : qo_tile_indices_offset(0), + qo_indptr_offset(0), + kv_indptr_offset(0), + qo_len_offset(0), + kv_len_offset(0), + head_indices_offset(0), + work_indptr_offset(0), + batch_indices_offset(0), + same_schedule_for_all_heads(false) {} + + // convert PrefillPlanSM90Info to std::vector + std::vector ToVector() const { + return {qo_tile_indices_offset, qo_indptr_offset, kv_indptr_offset, + qo_len_offset, kv_len_offset, head_indices_offset, + work_indptr_offset, batch_indices_offset, same_schedule_for_all_heads}; + } + + // From std::vector to PrefillPlanSM90Info + void FromVector(const std::vector& vec) { + if (vec.size() != 9) { + std::ostringstream err_msg; + err_msg << "PrefillPlanSM90Info::FromVector: vec.size() should be 9, but got " << vec.size(); + FLASHINFER_ERROR(err_msg.str()); + } + qo_tile_indices_offset = vec[0]; + qo_indptr_offset = vec[1]; + kv_indptr_offset = vec[2]; + qo_len_offset = vec[3]; + kv_len_offset = vec[4]; + head_indices_offset = vec[5]; + work_indptr_offset = vec[6]; + batch_indices_offset = vec[7]; + same_schedule_for_all_heads = vec[8]; + } +}; + +template +inline cudaError_t PrefillSM90Plan( + void* float_buffer, size_t float_workspace_size_in_bytes, void* int_buffer, + void* page_locked_int_buffer, size_t int_workspace_size_in_bytes, + PrefillPlanSM90Info& plan_info, IdType* qo_indptr_h, IdType* kv_indptr_h, IdType* kv_len_arr_h, + uint32_t total_num_rows, uint32_t batch_size, uint32_t num_qo_heads, uint32_t num_kv_heads, + uint32_t head_dim_qk, uint32_t head_dim_vo, uint32_t page_size, bool causal, + bool enable_cuda_graph, uint32_t sizeof_dtype_o, cudaStream_t stream) { + if (num_qo_heads % num_kv_heads != 0) { + std::ostringstream err_msg; + err_msg << "num_qo_heads " << num_qo_heads << " should be divisible by num_kv_heads " + << num_kv_heads; + FLASHINFER_ERROR(err_msg.str()); + } + + std::vector> idx_qo_kv_len_vec; + for (uint32_t i = 0; i < batch_size; ++i) { + int qo_len = qo_indptr_h[i + 1] - qo_indptr_h[i]; + int kv_len = kv_len_arr_h[i]; + if (kv_len < 0) { + std::ostringstream err_msg; + err_msg << "kv_len[" << i << "]" << kv_len << " should be non-negative"; + FLASHINFER_ERROR(err_msg.str()); + } + if (qo_len < 0) { + std::ostringstream err_msg; + err_msg << "qo_indptr[" << i + 1 << "]" << qo_indptr_h[i + 1] << " - qo_indptr[" << i << "]" + << qo_indptr_h[i] << " should be non-negative"; + FLASHINFER_ERROR(err_msg.str()); + } + idx_qo_kv_len_vec.push_back({i, qo_len, kv_len}); + } + + std::sort(idx_qo_kv_len_vec.begin(), idx_qo_kv_len_vec.end(), + [](const auto& a, const auto& b) { return std::get<2>(a) > std::get<2>(b); }); + int cta_tile_q = 128; + if (head_dim_vo == 64) { + cta_tile_q = 192; + } + + int device = 0; + FLASHINFER_CUDA_CALL(cudaGetDevice(&device)); + int num_sm90_ctas = 0; + FLASHINFER_CUDA_CALL( + cudaDeviceGetAttribute(&num_sm90_ctas, cudaDevAttrMultiProcessorCount, device)); + + MinHeap cta_cost_heap(num_sm90_ctas); + std::vector> cta_qo_tile_indices(num_sm90_ctas, std::vector()), + cta_qo_indptr(num_sm90_ctas, std::vector()), + cta_kv_indptr(num_sm90_ctas, std::vector()), + cta_qo_len(num_sm90_ctas, std::vector()), + cta_kv_len(num_sm90_ctas, std::vector()), + cta_head_indices(num_sm90_ctas, std::vector()), + cta_batch_indices(num_sm90_ctas, std::vector()); + + int max_num_works_per_head = ceil_div(total_num_rows, cta_tile_q) + batch_size - 1; + plan_info.same_schedule_for_all_heads = max_num_works_per_head > 4096; + + for (int qo_head_idx = 0; + qo_head_idx < (plan_info.same_schedule_for_all_heads ? 1 : num_qo_heads); ++qo_head_idx) { + for (auto& [i, qo_len, kv_len] : idx_qo_kv_len_vec) { + int num_qo_tiles = ceil_div(qo_len, cta_tile_q); + for (int qo_tile_idx = num_qo_tiles - 1; qo_tile_idx >= 0; --qo_tile_idx) { + auto [cta_idx, accum_cost] = cta_cost_heap.pop(); + // NOTE(Zihao): our current FA3 implementation do not fuse query and group heads + // so the group_size in cost_function is always 1 + int effective_kv_len = + causal ? packed_causal_kv_end(qo_len, kv_len, qo_tile_idx, cta_tile_q, num_qo_tiles, 1) + : kv_len; + cta_cost_heap.insert({cta_idx, accum_cost + cost_function(cta_tile_q, effective_kv_len)}); + cta_qo_tile_indices[cta_idx].push_back(qo_tile_idx); + cta_qo_indptr[cta_idx].push_back(qo_indptr_h[i]); + cta_qo_len[cta_idx].push_back(qo_len); + cta_kv_indptr[cta_idx].push_back(kv_indptr_h[i]); + cta_kv_len[cta_idx].push_back(kv_len); + cta_head_indices[cta_idx].push_back(qo_head_idx); + cta_batch_indices[cta_idx].push_back(i); + } + } + } + + std::vector work_indptr_vec(num_sm90_ctas + 1, 0); + for (uint32_t i = 0; i < num_sm90_ctas; ++i) { + work_indptr_vec[i + 1] = work_indptr_vec[i] + cta_qo_tile_indices[i].size(); + } + int total_num_works = work_indptr_vec.back(); + auto qo_tile_indices_vec = flatten(cta_qo_tile_indices, total_num_works); + auto qo_indptr_vec = flatten(cta_qo_indptr, total_num_works); + auto kv_indptr_vec = flatten(cta_kv_indptr, total_num_works); + auto qo_len_vec = flatten(cta_qo_len, total_num_works); + auto kv_len_vec = flatten(cta_kv_len, total_num_works); + auto head_indices_vec = flatten(cta_head_indices, total_num_works); + auto batch_indices_vec = flatten(cta_batch_indices, total_num_works); + + AlignedAllocator int_allocator(int_buffer, int_workspace_size_in_bytes); + int max_total_num_works; + + if (enable_cuda_graph) { + max_total_num_works = plan_info.same_schedule_for_all_heads + ? max_num_works_per_head + : max_num_works_per_head * num_qo_heads; + } else { + max_total_num_works = total_num_works; + } + + plan_info.qo_tile_indices_offset = int_allocator.aligned_alloc_offset( + sizeof(IdType) * max_total_num_works, 16, "batch_prefill_sm90_qo_tile_indices"); + plan_info.qo_indptr_offset = int_allocator.aligned_alloc_offset( + sizeof(IdType) * max_total_num_works, 16, "batch_prefill_sm90_qo_offset"); + plan_info.kv_indptr_offset = int_allocator.aligned_alloc_offset( + sizeof(IdType) * max_total_num_works, 16, "batch_prefill_sm90_kv_offset"); + plan_info.qo_len_offset = int_allocator.aligned_alloc_offset(sizeof(IdType) * max_total_num_works, + 16, "batch_prefill_sm90_qo_len"); + plan_info.kv_len_offset = int_allocator.aligned_alloc_offset(sizeof(IdType) * max_total_num_works, + 16, "batch_prefill_sm90_kv_len"); + plan_info.head_indices_offset = int_allocator.aligned_alloc_offset( + sizeof(IdType) * max_total_num_works, 16, "batch_prefill_sm90_head_indices"); + plan_info.work_indptr_offset = int_allocator.aligned_alloc_offset( + sizeof(IdType) * (num_sm90_ctas + 1), 16, "batch_prefill_sm90_work_indptr"); + plan_info.batch_indices_offset = int_allocator.aligned_alloc_offset( + sizeof(IdType) * max_total_num_works, 16, "batch_prefill_sm90_batch_indices"); + + IdType* qo_tile_indices_h = + GetPtrFromBaseOffset(page_locked_int_buffer, plan_info.qo_tile_indices_offset); + IdType* qo_offset_h = + GetPtrFromBaseOffset(page_locked_int_buffer, plan_info.qo_indptr_offset); + IdType* kv_offset_h = + GetPtrFromBaseOffset(page_locked_int_buffer, plan_info.kv_indptr_offset); + IdType* qo_len_h = GetPtrFromBaseOffset(page_locked_int_buffer, plan_info.qo_len_offset); + IdType* kv_len_h = GetPtrFromBaseOffset(page_locked_int_buffer, plan_info.kv_len_offset); + IdType* head_indices_h = + GetPtrFromBaseOffset(page_locked_int_buffer, plan_info.head_indices_offset); + IdType* work_indptr_h = + GetPtrFromBaseOffset(page_locked_int_buffer, plan_info.work_indptr_offset); + IdType* batch_indices_h = + GetPtrFromBaseOffset(page_locked_int_buffer, plan_info.batch_indices_offset); + + std::copy(qo_tile_indices_vec.begin(), qo_tile_indices_vec.end(), qo_tile_indices_h); + std::copy(qo_indptr_vec.begin(), qo_indptr_vec.end(), qo_offset_h); + std::copy(kv_indptr_vec.begin(), kv_indptr_vec.end(), kv_offset_h); + std::copy(qo_len_vec.begin(), qo_len_vec.end(), qo_len_h); + std::copy(kv_len_vec.begin(), kv_len_vec.end(), kv_len_h); + std::copy(head_indices_vec.begin(), head_indices_vec.end(), head_indices_h); + std::copy(work_indptr_vec.begin(), work_indptr_vec.end(), work_indptr_h); + std::copy(batch_indices_vec.begin(), batch_indices_vec.end(), batch_indices_h); + + size_t num_bytes_to_copy = int_allocator.num_allocated_bytes(); + FLASHINFER_CUDA_CALL(cudaMemcpyAsync(int_buffer, page_locked_int_buffer, num_bytes_to_copy, + cudaMemcpyHostToDevice, stream)); + return cudaSuccess; +} + +template +struct HolisticPlanInfo { + int64_t num_blks_x; + int64_t num_blks_y; + struct { + int64_t q_indptr_offset; + int64_t kv_indptr_offset; + int64_t partial_indptr_offset; + int64_t q_len_offset; + int64_t kv_len_offset; + int64_t q_start_offset; + int64_t kv_start_offset; + int64_t kv_end_offset; + int64_t kv_head_idx_offset; + int64_t work_indptr_offset; + int64_t len_kv_chunk_offset; + } tasks[NUM_TASKS]; + + int64_t partial_o_offset; + int64_t partial_lse_offset; + int64_t merge_indptr_offset; + int64_t merge_o_indices_offset; + int64_t num_qo_len_offset; + + static constexpr uint32_t NUM_TASK_ARGS = 11; + static constexpr uint32_t NUM_SHARED_ARGS = 7; + + std::vector ToVector() const { + std::vector vec; + vec.push_back(num_blks_x); + vec.push_back(num_blks_y); + for (uint32_t i = 0; i < NUM_TASKS; ++i) { + vec.push_back(tasks[i].q_indptr_offset); + vec.push_back(tasks[i].kv_indptr_offset); + vec.push_back(tasks[i].partial_indptr_offset); + vec.push_back(tasks[i].q_len_offset); + vec.push_back(tasks[i].kv_len_offset); + vec.push_back(tasks[i].q_start_offset); + vec.push_back(tasks[i].kv_start_offset); + vec.push_back(tasks[i].kv_end_offset); + vec.push_back(tasks[i].kv_head_idx_offset); + vec.push_back(tasks[i].work_indptr_offset); + vec.push_back(tasks[i].len_kv_chunk_offset); + } + vec.push_back(partial_o_offset); + vec.push_back(partial_lse_offset); + vec.push_back(merge_indptr_offset); + vec.push_back(merge_o_indices_offset); + vec.push_back(num_qo_len_offset); + return vec; + } + + void FromVector(const std::vector& vec) { + if (vec.size() != NUM_SHARED_ARGS + NUM_TASKS * NUM_TASK_ARGS) { + std::ostringstream err_msg; + err_msg << "HolisticPlanInfo::FromVector: vec.size() should be " + << NUM_SHARED_ARGS + NUM_TASKS * NUM_TASK_ARGS << ", but got " << vec.size(); + FLASHINFER_ERROR(err_msg.str()); + } + num_blks_x = vec[0]; + num_blks_y = vec[1]; + for (uint32_t i = 0; i < NUM_TASKS; ++i) { + tasks[i].q_indptr_offset = vec[2 + i * NUM_TASK_ARGS + 0]; + tasks[i].kv_indptr_offset = vec[2 + i * NUM_TASK_ARGS + 1]; + tasks[i].partial_indptr_offset = vec[2 + i * NUM_TASK_ARGS + 2]; + tasks[i].q_len_offset = vec[2 + i * NUM_TASK_ARGS + 3]; + tasks[i].kv_len_offset = vec[2 + i * NUM_TASK_ARGS + 4]; + tasks[i].q_start_offset = vec[2 + i * NUM_TASK_ARGS + 5]; + tasks[i].kv_start_offset = vec[2 + i * NUM_TASK_ARGS + 6]; + tasks[i].kv_end_offset = vec[2 + i * NUM_TASK_ARGS + 7]; + tasks[i].kv_head_idx_offset = vec[2 + i * NUM_TASK_ARGS + 8]; + tasks[i].work_indptr_offset = vec[2 + i * NUM_TASK_ARGS + 9]; + tasks[i].len_kv_chunk_offset = vec[2 + i * NUM_TASK_ARGS + 10]; + } + partial_o_offset = vec[2 + NUM_TASKS * NUM_TASK_ARGS]; + partial_lse_offset = vec[3 + NUM_TASKS * NUM_TASK_ARGS]; + merge_indptr_offset = vec[4 + NUM_TASKS * NUM_TASK_ARGS]; + merge_o_indices_offset = vec[5 + NUM_TASKS * NUM_TASK_ARGS]; + num_qo_len_offset = vec[6 + NUM_TASKS * NUM_TASK_ARGS]; + } +}; + +template +inline cudaError_t TwoStageHolisticPlan(void* float_buffer, size_t float_workspace_size_in_bytes, + void* int_buffer, void* page_locked_int_buffer, + size_t int_workspace_size_in_bytes, + HolisticPlanInfo<2>& plan_info, IdType* qo_indptr_h, + IdType* kv_indptr_h, IdType* kv_len_arr_h, + uint32_t batch_size, uint32_t num_qo_heads, + uint32_t num_kv_heads, uint32_t head_dim, bool causal, + cudaStream_t stream) { + constexpr uint32_t NUM_TASKS = 2; + const uint32_t CTA_TILE_Q_SIZES[NUM_TASKS] = {128, 16}; + int num_sm = 0; + int dev_id = 0; + + uint32_t gqa_group_size = num_qo_heads / num_kv_heads; + FLASHINFER_CUDA_CALL(cudaGetDevice(&dev_id)); + FLASHINFER_CUDA_CALL(cudaDeviceGetAttribute(&num_sm, cudaDevAttrMultiProcessorCount, dev_id)); + + if (head_dim >= 256) { + // NOTE (Yilong): optimize this code path + // constraint gridDim due to cooperative group + num_sm *= 1; + } else { + // NOTE(Zihao): two cta per sm + num_sm *= 2; + } + + // step 0. determine the number of blocks in x and y dimensions + std::vector> idx_qo_kv_len_vec[NUM_TASKS]; + for (uint32_t i = 0; i < batch_size; ++i) { + if (qo_indptr_h[i + 1] - qo_indptr_h[i] < 0) { + std::ostringstream err_msg; + err_msg << "qo_indptr[" << i + 1 << "]" << qo_indptr_h[i + 1] << " - qo_indptr[" << i << "]" + << qo_indptr_h[i] << " should be non-negative"; + FLASHINFER_ERROR(err_msg.str()); + } + + int qo_len = qo_indptr_h[i + 1] - qo_indptr_h[i]; + int packed_qo_len = qo_len * gqa_group_size; + int kv_len = kv_len_arr_h[i]; + + if (packed_qo_len > CTA_TILE_Q_SIZES[1]) { + idx_qo_kv_len_vec[0].push_back({i, qo_len, kv_len}); + } else { + idx_qo_kv_len_vec[1].push_back({i, qo_len, kv_len}); + } + } + + int cluster_size = 1; + int num_clusters = num_sm / cluster_size; + plan_info.num_blks_x = cluster_size; + plan_info.num_blks_y = num_clusters; + + auto f = [](int x) { + if (x <= 128) { + // This aligns with CTA_TILE_KV in persistent mainloop + // NOTE (Yilong): Optimize here for smaller batch/seqlen scenarios + return 128; + } + return ceil_div(x, 256) * 256; + }; + + MinHeap cluster_cost_heap(num_clusters); + AlignedAllocator int_allocator(int_buffer, int_workspace_size_in_bytes); + + // NOTE(Zihao): adjust it later + const int max_total_num_works = 65536; + const int max_num_kv_splits = + 4 * num_clusters * cluster_size * (CTA_TILE_Q_SIZES[0] + CTA_TILE_Q_SIZES[1]); + + // calculate kv_len_limit first, considering all workloads + int64_t total_kv_lens = 0; + for (uint32_t task = 0; task < NUM_TASKS; ++task) { + int cluster_tile_q = CTA_TILE_Q_SIZES[task] * cluster_size; + for (auto& [_, qo_len, kv_len] : idx_qo_kv_len_vec[task]) { + int packed_qo_len = qo_len * gqa_group_size; + int num_qo_tiles = ceil_div(packed_qo_len, cluster_tile_q); + for (int qo_tile_idx = num_qo_tiles - 1; qo_tile_idx >= 0; --qo_tile_idx) { + int effective_kv_len = + causal ? packed_causal_kv_end(qo_len, kv_len, qo_tile_idx, cluster_tile_q, num_qo_tiles, + gqa_group_size) + : kv_len; + total_kv_lens += effective_kv_len; + } + } + } + + // used for remapping the output offsets + // layout [packed_qo_len x num_kv_tiles, num_kv_heads, head_dim] + int partial_o_nnz = 0; + std::vector merge_indptr, merge_o_indices, num_expand_qo_len_vec; + merge_indptr.push_back(partial_o_nnz); + for (uint32_t task = 0; task < NUM_TASKS; ++task) { + int cluster_tile_q = CTA_TILE_Q_SIZES[task] * cluster_size; + int kv_len_limit = 0; + if (cluster_tile_q >= 64) { + // chunked-prefill workloads are much more expensive than decode + // so we use a smaller kv_len_limit for chunked-prefill workloads + kv_len_limit = f(std::max(ceil_div(total_kv_lens, num_clusters), 1L)); + } else { + kv_len_limit = f(std::max(ceil_div(total_kv_lens * num_kv_heads, num_clusters), 1L)); + } + + std::vector> cluster_q_indptr(num_clusters, std::vector()), + cluster_kv_indptr(num_clusters, std::vector()), + cluster_q_len(num_clusters, std::vector()), + cluster_kv_len(num_clusters, std::vector()), + cluster_q_start(num_clusters, std::vector()), + cluster_kv_start(num_clusters, std::vector()), + cluster_kv_end(num_clusters, std::vector()), + cluster_kv_head_idx(num_clusters, std::vector()), + cluster_partial_indptr(num_clusters, std::vector()), + cluster_len_kv_chunk(num_clusters, std::vector()); + + for (auto& [i, qo_len, kv_len] : idx_qo_kv_len_vec[task]) { + int packed_qo_len = qo_len * gqa_group_size; + int num_qo_tiles = ceil_div(packed_qo_len, cluster_tile_q); + // NOTE (Yilong): this ordering correspoinds to the layout of reduction kernel + for (int qo_tile_idx = 0; qo_tile_idx < num_qo_tiles; ++qo_tile_idx) { + int remaining_len = causal + ? packed_causal_kv_end(qo_len, kv_len, qo_tile_idx, cluster_tile_q, + num_qo_tiles, gqa_group_size) + : kv_len; + int kv_start = 0; + bool split_kv = remaining_len > kv_len_limit; + int num_kv_tiles = split_kv ? ceil_div(remaining_len, kv_len_limit) : 1; + int row_tile_size = std::min(cluster_tile_q, packed_qo_len - qo_tile_idx * cluster_tile_q); + bool zero_kv_len = (remaining_len == 0); + while (remaining_len > 0 || zero_kv_len) { + int actual_len = std::min(remaining_len, kv_len_limit); + for (uint32_t kv_head_idx = 0; kv_head_idx < num_kv_heads; ++kv_head_idx) { + auto [cluster_idx, accum_cost] = cluster_cost_heap.pop(); + cluster_cost_heap.insert( + {cluster_idx, accum_cost + cost_function(cluster_tile_q, actual_len)}); + cluster_q_len[cluster_idx].push_back(qo_len); + cluster_kv_len[cluster_idx].push_back(kv_len); + cluster_q_indptr[cluster_idx].push_back(qo_indptr_h[i]); + cluster_kv_indptr[cluster_idx].push_back(kv_indptr_h[i]); + + // use kv_chunk to rematerize num_kv_tiles and kv_tile_idx + cluster_len_kv_chunk[cluster_idx].push_back(kv_len_limit); + cluster_partial_indptr[cluster_idx].push_back(partial_o_nnz); + + cluster_q_start[cluster_idx].push_back(qo_tile_idx * cluster_tile_q); + cluster_kv_start[cluster_idx].push_back(kv_start); + cluster_kv_end[cluster_idx].push_back(kv_start + actual_len); + cluster_kv_head_idx[cluster_idx].push_back(kv_head_idx); + } + remaining_len -= actual_len; + zero_kv_len = (remaining_len == 0); + kv_start += actual_len; + if (zero_kv_len) { + break; + } + } + if (split_kv) { + // non-split kv is directly written through + for (int row = 0; row < row_tile_size; ++row) { + merge_indptr.push_back(merge_indptr.back() + num_kv_tiles); + merge_o_indices.push_back(qo_indptr_h[i] + + (qo_tile_idx * cluster_tile_q + row) / gqa_group_size); + } + partial_o_nnz += row_tile_size * num_kv_tiles; + } + } + } + + std::vector work_indptr_vec(num_clusters + 1, 0); + for (uint32_t i = 0; i < num_clusters; ++i) { + work_indptr_vec[i + 1] = work_indptr_vec[i] + cluster_q_indptr[i].size(); + } + int total_num_works = work_indptr_vec.back(); + if (total_num_works > max_total_num_works) { + std::ostringstream err_msg; + err_msg << "total_num_works (#q tiles * #kv tiles) " << total_num_works + << " exceeds max_total_num_works " << max_total_num_works; + FLASHINFER_ERROR(err_msg.str()); + } + auto q_indptr_vec = flatten(cluster_q_indptr, total_num_works); + auto kv_indptr_vec = flatten(cluster_kv_indptr, total_num_works); + auto partial_indptr_vec = flatten(cluster_partial_indptr, total_num_works); + auto q_len_vec = flatten(cluster_q_len, total_num_works); + auto kv_len_vec = flatten(cluster_kv_len, total_num_works); + auto q_start_vec = flatten(cluster_q_start, total_num_works); + auto kv_start_vec = flatten(cluster_kv_start, total_num_works); + auto kv_end_vec = flatten(cluster_kv_end, total_num_works); + auto kv_head_idx_vec = flatten(cluster_kv_head_idx, total_num_works); + auto len_kv_chunk_vec = flatten(cluster_len_kv_chunk, total_num_works); + + plan_info.tasks[task].q_indptr_offset = + int_allocator.aligned_alloc_offset(sizeof(IdType) * max_total_num_works, 16, "q_indptr"); + plan_info.tasks[task].kv_indptr_offset = + int_allocator.aligned_alloc_offset(sizeof(IdType) * max_total_num_works, 16, "kv_indptr"); + plan_info.tasks[task].partial_indptr_offset = int_allocator.aligned_alloc_offset( + sizeof(IdType) * max_total_num_works, 16, "partial_indptr"); + plan_info.tasks[task].q_len_offset = + int_allocator.aligned_alloc_offset(sizeof(IdType) * max_total_num_works, 16, "q_len"); + plan_info.tasks[task].kv_len_offset = + int_allocator.aligned_alloc_offset(sizeof(IdType) * max_total_num_works, 16, "kv_len"); + plan_info.tasks[task].q_start_offset = + int_allocator.aligned_alloc_offset(sizeof(IdType) * max_total_num_works, 16, "q_start"); + plan_info.tasks[task].kv_start_offset = + int_allocator.aligned_alloc_offset(sizeof(IdType) * max_total_num_works, 16, "kv_start"); + plan_info.tasks[task].kv_end_offset = + int_allocator.aligned_alloc_offset(sizeof(IdType) * max_total_num_works, 16, "kv_end"); + plan_info.tasks[task].kv_head_idx_offset = + int_allocator.aligned_alloc_offset(sizeof(IdType) * max_total_num_works, 16, "kv_head_idx"); + plan_info.tasks[task].work_indptr_offset = + int_allocator.aligned_alloc_offset(sizeof(IdType) * max_total_num_works, 16, "work_indptr"); + plan_info.tasks[task].len_kv_chunk_offset = int_allocator.aligned_alloc_offset( + sizeof(IdType) * max_total_num_works, 16, "len_kv_chunk"); + + CopyToPageLockedBuffer(page_locked_int_buffer, plan_info.tasks[task].q_indptr_offset, + q_indptr_vec); + CopyToPageLockedBuffer(page_locked_int_buffer, plan_info.tasks[task].kv_indptr_offset, + kv_indptr_vec); + CopyToPageLockedBuffer(page_locked_int_buffer, plan_info.tasks[task].partial_indptr_offset, + partial_indptr_vec); + CopyToPageLockedBuffer(page_locked_int_buffer, plan_info.tasks[task].q_len_offset, q_len_vec); + CopyToPageLockedBuffer(page_locked_int_buffer, plan_info.tasks[task].kv_len_offset, kv_len_vec); + CopyToPageLockedBuffer(page_locked_int_buffer, plan_info.tasks[task].q_start_offset, + q_start_vec); + CopyToPageLockedBuffer(page_locked_int_buffer, plan_info.tasks[task].kv_start_offset, + kv_start_vec); + CopyToPageLockedBuffer(page_locked_int_buffer, plan_info.tasks[task].kv_end_offset, kv_end_vec); + CopyToPageLockedBuffer(page_locked_int_buffer, plan_info.tasks[task].kv_head_idx_offset, + kv_head_idx_vec); + CopyToPageLockedBuffer(page_locked_int_buffer, plan_info.tasks[task].work_indptr_offset, + work_indptr_vec); + CopyToPageLockedBuffer(page_locked_int_buffer, plan_info.tasks[task].len_kv_chunk_offset, + len_kv_chunk_vec); + } + + if (merge_indptr.size() > max_num_kv_splits) { + std::ostringstream err_msg; + err_msg << "Number of kv splits " << merge_indptr.size() << " exceeds max buffer size " + << max_num_kv_splits << ". Please increase the threshold."; + FLASHINFER_ERROR(err_msg.str()); + } + + // update num_qo_len_vec + num_expand_qo_len_vec.push_back(merge_indptr.size() - 1); + // allocate buffer for state merge function + plan_info.merge_indptr_offset = + int_allocator.aligned_alloc_offset(sizeof(IdType) * max_num_kv_splits, 16, "merge_indptr"); + plan_info.merge_o_indices_offset = + int_allocator.aligned_alloc_offset(sizeof(IdType) * max_num_kv_splits, 16, "merge_o_indices"); + plan_info.num_qo_len_offset = + int_allocator.aligned_alloc_offset(sizeof(IdType), 16, "num_qo_len_offset"); + // copy data to paged cpu buffer + CopyToPageLockedBuffer(page_locked_int_buffer, plan_info.merge_indptr_offset, merge_indptr); + CopyToPageLockedBuffer(page_locked_int_buffer, plan_info.merge_o_indices_offset, merge_o_indices); + CopyToPageLockedBuffer(page_locked_int_buffer, plan_info.num_qo_len_offset, + num_expand_qo_len_vec); + + size_t num_bytes_to_copy = int_allocator.num_allocated_bytes(); + FLASHINFER_CUDA_CALL(cudaMemcpyAsync(int_buffer, page_locked_int_buffer, num_bytes_to_copy, + cudaMemcpyHostToDevice, stream)); + constexpr size_t sizeof_dtype_o = 2; // NOTE (Yilong): assume fp16 + + // Note(Yilong): adjust it later + AlignedAllocator float_allocator(float_buffer, float_workspace_size_in_bytes); + plan_info.partial_o_offset = float_allocator.aligned_alloc_offset( + 2 * max_num_kv_splits * sizeof_dtype_o * head_dim, 16, "holistic_partial_o"); + plan_info.partial_lse_offset = float_allocator.aligned_alloc_offset( + 2 * max_num_kv_splits * sizeof(float), 16, "holistic_partial_lse"); + + return cudaSuccess; +} + +struct MLAPlanInfo { + int64_t num_blks_x; + int64_t num_blks_y; + int64_t q_indptr_offset; + int64_t kv_indptr_offset; + int64_t partial_indptr_offset; + int64_t merge_packed_offset_start_offset; + int64_t merge_packed_offset_end_offset; + int64_t merge_partial_packed_offset_start_offset; + int64_t merge_partial_packed_offset_end_offset; + int64_t merge_partial_stride_offset; + int64_t q_len_offset; + int64_t kv_len_offset; + int64_t q_start_offset; + int64_t kv_start_offset; + int64_t kv_end_offset; + int64_t work_indptr_offset; + int64_t partial_o_offset; + int64_t partial_lse_offset; + + std::vector ToVector() const { + return {num_blks_x, + num_blks_y, + q_indptr_offset, + kv_indptr_offset, + partial_indptr_offset, + merge_packed_offset_start_offset, + merge_packed_offset_end_offset, + merge_partial_packed_offset_start_offset, + merge_partial_packed_offset_end_offset, + merge_partial_stride_offset, + q_len_offset, + kv_len_offset, + q_start_offset, + kv_start_offset, + kv_end_offset, + work_indptr_offset, + partial_o_offset, + partial_lse_offset}; + } + + void FromVector(const std::vector& vec) { + if (vec.size() != 18) { + std::ostringstream err_msg; + err_msg << "MLAPlanInfo::FromVector: vec.size() should be 18, but got " << vec.size(); + FLASHINFER_ERROR(err_msg.str()); + } + num_blks_x = vec[0]; + num_blks_y = vec[1]; + q_indptr_offset = vec[2]; + kv_indptr_offset = vec[3]; + partial_indptr_offset = vec[4]; + merge_packed_offset_start_offset = vec[5]; + merge_packed_offset_end_offset = vec[6]; + merge_partial_packed_offset_start_offset = vec[7]; + merge_partial_packed_offset_end_offset = vec[8]; + merge_partial_stride_offset = vec[9]; + q_len_offset = vec[10]; + kv_len_offset = vec[11]; + q_start_offset = vec[12]; + kv_start_offset = vec[13]; + kv_end_offset = vec[14]; + work_indptr_offset = vec[15]; + partial_o_offset = vec[16]; + partial_lse_offset = vec[17]; + } +}; + +template +inline cudaError_t MLAPlan(void* float_buffer, size_t float_workspace_size_in_bytes, + void* int_buffer, void* page_locked_int_buffer, + size_t int_workspace_size_in_bytes, MLAPlanInfo& plan_info, + IdType* qo_indptr_h, IdType* kv_indptr_h, IdType* kv_len_arr_h, + uint32_t batch_size, uint32_t num_heads, uint32_t head_dim_o, + bool causal, cudaStream_t stream) { + int num_sm = 0; + int dev_id = 0; + FLASHINFER_CUDA_CALL(cudaGetDevice(&dev_id)); + FLASHINFER_CUDA_CALL(cudaDeviceGetAttribute(&num_sm, cudaDevAttrMultiProcessorCount, dev_id)); + + // step 0. determine the number of blocks in x and y dimensions + int accum_packed_qo_len = 0; + std::vector> idx_qo_kv_len_vec; + for (uint32_t i = 0; i < batch_size; ++i) { + if (qo_indptr_h[i + 1] - qo_indptr_h[i] < 0) { + std::ostringstream err_msg; + err_msg << "qo_indptr[" << i + 1 << "]" << qo_indptr_h[i + 1] << " - qo_indptr[" << i << "]" + << qo_indptr_h[i] << " should be non-negative"; + FLASHINFER_ERROR(err_msg.str()); + } + + int qo_len = qo_indptr_h[i + 1] - qo_indptr_h[i]; + int packed_qo_len = qo_len * num_heads; + accum_packed_qo_len += packed_qo_len; + + int kv_len = kv_len_arr_h[i]; + idx_qo_kv_len_vec.push_back({i, qo_len, kv_len}); + } + int avg_packed_qo_len = accum_packed_qo_len / batch_size; + + int cluster_size; + if (avg_packed_qo_len > 64) { + cluster_size = 2; // two ctas in a cluster + } else { + cluster_size = 1; // one cta in a cluster + } + uint32_t num_clusters = num_sm / cluster_size; + plan_info.num_blks_x = cluster_size; + plan_info.num_blks_y = num_clusters; + const int cta_tile_q = 64; + int cluster_tile_q = cluster_size * cta_tile_q; + + int64_t total_kv_lens = 0; + for (auto& [_, qo_len, kv_len] : idx_qo_kv_len_vec) { + int packed_qo_len = qo_len * num_heads; + int num_qo_tiles = ceil_div(packed_qo_len, cluster_tile_q); + for (int qo_tile_idx = num_qo_tiles - 1; qo_tile_idx >= 0; --qo_tile_idx) { + int effective_kv_len = causal ? packed_causal_kv_end(qo_len, kv_len, qo_tile_idx, + cluster_tile_q, num_qo_tiles, num_heads) + : kv_len; + total_kv_lens += effective_kv_len; + } + } + + auto f = [](int x) { + if (x <= 8) { + return 32; + } else if (x <= 16) { + return 64; + } else if (x <= 32) { + return 128; + } else if (x <= 64) { + return 192; + } + return ceil_div(x, 256) * 256; + }; + + int kv_len_limit = f(std::max(ceil_div(total_kv_lens, num_clusters), 1L)); + + // step 1. load-balancing scheduling algorithm + MinHeap cluster_cost_heap(num_clusters); + std::vector> cluster_q_indptr(num_clusters, std::vector()), + cluster_kv_indptr(num_clusters, std::vector()), + cluster_q_len(num_clusters, std::vector()), + cluster_kv_len(num_clusters, std::vector()), + cluster_q_start(num_clusters, std::vector()), + cluster_kv_start(num_clusters, std::vector()), + cluster_kv_end(num_clusters, std::vector()), + cluster_partial_indptr(num_clusters, std::vector()); + + std::vector merge_packed_offset_start(num_sm, 0), merge_packed_offset_end(num_sm, 0), + merge_partial_packed_offset_start(num_sm, 0), merge_partial_packed_offset_end(num_sm, 0), + merge_partial_stride(num_sm, 0); + + int merge_cta_counter = 0; + int partial_o_nnz = 0; + + for (auto& [i, qo_len, kv_len] : idx_qo_kv_len_vec) { + int packed_qo_len = qo_len * num_heads; + int num_qo_tiles = ceil_div(packed_qo_len, cluster_tile_q); + for (int qo_tile_idx = num_qo_tiles - 1; qo_tile_idx >= 0; --qo_tile_idx) { + int remaining_len = causal ? packed_causal_kv_end(qo_len, kv_len, qo_tile_idx, cluster_tile_q, + num_qo_tiles, num_heads) + : kv_len; + int kv_start = 0; + bool split_kv = remaining_len > kv_len_limit; + int row_tile_size = std::min(cluster_tile_q, packed_qo_len - qo_tile_idx * cluster_tile_q); + if (split_kv) { + /* + * Proof(Zihao): merge_cta_counter <= num_sm (num_sm == num_clusters * cluster_size) + * + * Precondition: + * 1. kv_len_limit * num_clusters >= total_kv_lens == sum(remaining_len) + * 2. num_qo_chunks <= max((remaining_len * cluster_size) // kv_len_limit, 1) + * 3. num_qo_tiles_requires_split <= num_clusters + + * Implication: + * 1. sum(num_qo_chunks) <= max(sum(remaining_len) * cluster_size / kv_len_limit, + num_qo_tiles_requires_split) + * 2. sum(num_qo_chunks) <= max(cluster_size * num_clusters, num_qo_tiles_requires_split) + */ + int num_qo_chunks = std::max(remaining_len * cluster_size / kv_len_limit, 1); + // row_chunk_size * num_qo_chunks >= row_tile_size + int row_chunk_size = ceil_div(row_tile_size, num_qo_chunks); + int current_q_tile_end = + std::min(cluster_tile_q, packed_qo_len - qo_tile_idx * cluster_tile_q); + for (int offset_start = 0; offset_start < row_tile_size; offset_start += row_chunk_size) { + merge_packed_offset_start[merge_cta_counter] = + qo_indptr_h[i] * num_heads + qo_tile_idx * cluster_tile_q + offset_start; + merge_packed_offset_end[merge_cta_counter] = + qo_indptr_h[i] * num_heads + qo_tile_idx * cluster_tile_q + + std::min(offset_start + row_chunk_size, current_q_tile_end); + merge_partial_packed_offset_start[merge_cta_counter] = partial_o_nnz + offset_start; + merge_partial_packed_offset_end[merge_cta_counter] = + partial_o_nnz + ceil_div(remaining_len, kv_len_limit) * row_tile_size; + merge_partial_stride[merge_cta_counter] = row_tile_size; + merge_cta_counter++; + } + } + bool zero_kv_len = (remaining_len == 0); + while (remaining_len > 0 || zero_kv_len) { + auto [cluster_idx, accum_cost] = cluster_cost_heap.pop(); + int actual_len = std::min(remaining_len, kv_len_limit); + cluster_cost_heap.insert( + {cluster_idx, accum_cost + cost_function(cluster_tile_q, actual_len)}); + cluster_q_len[cluster_idx].push_back(qo_len); + cluster_kv_len[cluster_idx].push_back(kv_len); + cluster_q_indptr[cluster_idx].push_back(qo_indptr_h[i]); + cluster_kv_indptr[cluster_idx].push_back(kv_indptr_h[i]); + if (split_kv) { + cluster_partial_indptr[cluster_idx].push_back(partial_o_nnz); + partial_o_nnz += row_tile_size; + } else { + cluster_partial_indptr[cluster_idx].push_back(-1); + } + cluster_q_start[cluster_idx].push_back(qo_tile_idx * cluster_tile_q); + cluster_kv_start[cluster_idx].push_back(kv_start); + cluster_kv_end[cluster_idx].push_back(kv_start + actual_len); + remaining_len -= actual_len; + kv_start += actual_len; + if (zero_kv_len) break; + } + } + } + + FLASHINFER_CHECK(merge_cta_counter <= num_sm, + "Internal Error: merge_cta_counter should be less than or equal to num_sm, " + "please report this bug to the developers"); + + int max_total_num_works = 16384; // NOTE(Zihao): adjust it later + + std::vector work_indptr_vec(num_clusters + 1, 0); + for (uint32_t i = 0; i < num_clusters; ++i) { + work_indptr_vec[i + 1] = work_indptr_vec[i] + cluster_q_indptr[i].size(); + } + int total_num_works = work_indptr_vec.back(); + auto q_indptr_vec = flatten(cluster_q_indptr, total_num_works); + auto kv_indptr_vec = flatten(cluster_kv_indptr, total_num_works); + auto partial_indptr_vec = flatten(cluster_partial_indptr, total_num_works); + auto q_len_vec = flatten(cluster_q_len, total_num_works); + auto kv_len_vec = flatten(cluster_kv_len, total_num_works); + auto q_start_vec = flatten(cluster_q_start, total_num_works); + auto kv_start_vec = flatten(cluster_kv_start, total_num_works); + auto kv_end_vec = flatten(cluster_kv_end, total_num_works); + + AlignedAllocator int_allocator(int_buffer, int_workspace_size_in_bytes); + plan_info.q_indptr_offset = + int_allocator.aligned_alloc_offset(sizeof(IdType) * max_total_num_works, 16, "mla_q_indptr"); + plan_info.kv_indptr_offset = + int_allocator.aligned_alloc_offset(sizeof(IdType) * max_total_num_works, 16, "mla_kv_indptr"); + plan_info.partial_indptr_offset = int_allocator.aligned_alloc_offset( + sizeof(IdType) * max_total_num_works, 16, "mla_partial_indptr"); + plan_info.merge_packed_offset_start_offset = int_allocator.aligned_alloc_offset( + sizeof(IdType) * num_sm, 16, "mla_merge_packed_offset_start"); + plan_info.merge_packed_offset_end_offset = int_allocator.aligned_alloc_offset( + sizeof(IdType) * num_sm, 16, "mla_merge_packed_offset_end"); + plan_info.merge_partial_packed_offset_start_offset = int_allocator.aligned_alloc_offset( + sizeof(IdType) * num_sm, 16, "mla_merge_partial_packed_offset_start"); + plan_info.merge_partial_packed_offset_end_offset = int_allocator.aligned_alloc_offset( + sizeof(IdType) * num_sm, 16, "mla_merge_partial_packed_offset_end"); + plan_info.merge_partial_stride_offset = + int_allocator.aligned_alloc_offset(sizeof(IdType) * num_sm, 16, "mla_merge_partial_stride"); + plan_info.q_len_offset = + int_allocator.aligned_alloc_offset(sizeof(IdType) * max_total_num_works, 16, "mla_q_len"); + plan_info.kv_len_offset = + int_allocator.aligned_alloc_offset(sizeof(IdType) * max_total_num_works, 16, "mla_kv_len"); + plan_info.q_start_offset = + int_allocator.aligned_alloc_offset(sizeof(IdType) * max_total_num_works, 16, "mla_q_start"); + plan_info.kv_start_offset = + int_allocator.aligned_alloc_offset(sizeof(IdType) * max_total_num_works, 16, "mla_kv_start"); + plan_info.kv_end_offset = + int_allocator.aligned_alloc_offset(sizeof(IdType) * max_total_num_works, 16, "mla_kv_end"); + plan_info.work_indptr_offset = int_allocator.aligned_alloc_offset( + sizeof(IdType) * max_total_num_works, 16, "mla_work_indptr"); + + IdType* cluster_q_indptr_h = + GetPtrFromBaseOffset(page_locked_int_buffer, plan_info.q_indptr_offset); + IdType* cluster_kv_indptr_h = + GetPtrFromBaseOffset(page_locked_int_buffer, plan_info.kv_indptr_offset); + IdType* cluster_partial_indptr_h = + GetPtrFromBaseOffset(page_locked_int_buffer, plan_info.partial_indptr_offset); + IdType* cluster_merge_packed_offset_start_h = GetPtrFromBaseOffset( + page_locked_int_buffer, plan_info.merge_packed_offset_start_offset); + IdType* cluster_merge_packed_offset_end_h = GetPtrFromBaseOffset( + page_locked_int_buffer, plan_info.merge_packed_offset_end_offset); + IdType* cluster_merge_partial_packed_offset_start_h = GetPtrFromBaseOffset( + page_locked_int_buffer, plan_info.merge_partial_packed_offset_start_offset); + IdType* cluster_merge_partial_packed_offset_end_h = GetPtrFromBaseOffset( + page_locked_int_buffer, plan_info.merge_partial_packed_offset_end_offset); + IdType* cluster_merge_partial_stride_h = + GetPtrFromBaseOffset(page_locked_int_buffer, plan_info.merge_partial_stride_offset); + IdType* cluster_q_len_h = + GetPtrFromBaseOffset(page_locked_int_buffer, plan_info.q_len_offset); + IdType* cluster_kv_len_h = + GetPtrFromBaseOffset(page_locked_int_buffer, plan_info.kv_len_offset); + IdType* cluster_q_start_h = + GetPtrFromBaseOffset(page_locked_int_buffer, plan_info.q_start_offset); + IdType* cluster_kv_start_h = + GetPtrFromBaseOffset(page_locked_int_buffer, plan_info.kv_start_offset); + IdType* cluster_kv_end_h = + GetPtrFromBaseOffset(page_locked_int_buffer, plan_info.kv_end_offset); + IdType* cluster_work_indptr_h = + GetPtrFromBaseOffset(page_locked_int_buffer, plan_info.work_indptr_offset); + + std::copy(q_indptr_vec.begin(), q_indptr_vec.end(), cluster_q_indptr_h); + std::copy(kv_indptr_vec.begin(), kv_indptr_vec.end(), cluster_kv_indptr_h); + std::copy(partial_indptr_vec.begin(), partial_indptr_vec.end(), cluster_partial_indptr_h); + std::copy(merge_packed_offset_start.begin(), merge_packed_offset_start.end(), + cluster_merge_packed_offset_start_h); + std::copy(merge_packed_offset_end.begin(), merge_packed_offset_end.end(), + cluster_merge_packed_offset_end_h); + std::copy(merge_partial_packed_offset_start.begin(), merge_partial_packed_offset_start.end(), + cluster_merge_partial_packed_offset_start_h); + std::copy(merge_partial_packed_offset_end.begin(), merge_partial_packed_offset_end.end(), + cluster_merge_partial_packed_offset_end_h); + std::copy(merge_partial_stride.begin(), merge_partial_stride.end(), + cluster_merge_partial_stride_h); + std::copy(q_len_vec.begin(), q_len_vec.end(), cluster_q_len_h); + std::copy(kv_len_vec.begin(), kv_len_vec.end(), cluster_kv_len_h); + std::copy(q_start_vec.begin(), q_start_vec.end(), cluster_q_start_h); + std::copy(kv_start_vec.begin(), kv_start_vec.end(), cluster_kv_start_h); + std::copy(kv_end_vec.begin(), kv_end_vec.end(), cluster_kv_end_h); + std::copy(work_indptr_vec.begin(), work_indptr_vec.end(), cluster_work_indptr_h); + + size_t num_bytes_to_copy = int_allocator.num_allocated_bytes(); + FLASHINFER_CUDA_CALL(cudaMemcpyAsync(int_buffer, page_locked_int_buffer, num_bytes_to_copy, + cudaMemcpyHostToDevice, stream)); + + constexpr size_t sizeof_dtype_o = 2; + AlignedAllocator float_allocator(float_buffer, float_workspace_size_in_bytes); + plan_info.partial_o_offset = float_allocator.aligned_alloc_offset( + 2 * num_clusters * cluster_tile_q * sizeof_dtype_o * head_dim_o, 16, "mla_partial_o"); + plan_info.partial_lse_offset = float_allocator.aligned_alloc_offset( + 2 * num_clusters * cluster_tile_q * sizeof(float), 16, "mla_partial_lse"); + + return cudaSuccess; +} + +} // namespace flashinfer +#endif // FLASHINFER_ATTENTION_SCHEDULER_CUH_ diff --git a/include/flashinfer/attention/state.cuh b/include/flashinfer/attention/state.cuh new file mode 100644 index 0000000000000000000000000000000000000000..10789810120ebd7875ad69119788057785757e53 --- /dev/null +++ b/include/flashinfer/attention/state.cuh @@ -0,0 +1,83 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef FLASHINFER_STATE_CUH_ +#define FLASHINFER_STATE_CUH_ + +#include "../math.cuh" +#include "../vec_dtypes.cuh" + +namespace flashinfer { + +/*! + * \brief The flashattention state. + * \tparam vec_size The size of the vector used in o. + */ +template +struct state_t { + /* the weighted sum of v: exp(pre-softmax logit - m) * v / d */ + vec_t o; + /* maximum value of pre-softmax logits */ + float m; + /* sum of exp(pre-softmax logits - m) */ + float d; + + __device__ __forceinline__ void init() { + o.fill(0.f); + m = -math::inf; + d = 1.f; + } + + __device__ __forceinline__ state_t() { init(); } + + __device__ __forceinline__ float get_lse() const { return m + math::ptx_log2(d); } + + /*! + * \brief Merge the state with another state. + * \param other_m The maximum value of pre-softmax logits of the other state. + * \param other_d The sum of exp(pre-softmax logits - m) of the other state. + * \param other_o The weighted sum of v of the other state. + */ + __device__ __forceinline__ void merge(const vec_t& other_o, float other_m, + float other_d) { + float m_prev = m, d_prev = d; + m = max(m_prev, other_m); + d = d_prev * math::ptx_exp2(m_prev - m) + other_d * math::ptx_exp2(other_m - m); +#pragma unroll + for (size_t i = 0; i < vec_size; ++i) { + o[i] = o[i] * math::ptx_exp2(m_prev - m) + other_o[i] * math::ptx_exp2(other_m - m); + } + } + + /*! + * \brief Merge the state with another state. + * \param other The other state. + */ + __device__ __forceinline__ void merge(const state_t& other) { + merge(other.o, other.m, other.d); + } + + __device__ __forceinline__ void normalize() { + // only normalize by d when not normalized on the fly +#pragma unroll + for (size_t i = 0; i < vec_size; ++i) { + o[i] = __fdividef(o[i], d); + } + } +}; + +} // namespace flashinfer + +#endif // FLASHINFER_STATE_CUH_ diff --git a/include/flashinfer/attention/variant_helper.cuh b/include/flashinfer/attention/variant_helper.cuh new file mode 100644 index 0000000000000000000000000000000000000000..65a9f192492afdf0cac3dacfdb3a58182768c530 --- /dev/null +++ b/include/flashinfer/attention/variant_helper.cuh @@ -0,0 +1,89 @@ +/* + * Copyright (c) 2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef FLASHINFER_ATTENTION_VARIANT_HELPER_H +#define FLASHINFER_ATTENTION_VARIANT_HELPER_H + +#include + +#include + +namespace flashinfer { + +#define REGISTER_QUERY_TRANSFORM(params, q, ...) \ + template \ + __device__ __forceinline__ T QueryTransform(const Params& params, void* q_smem) { \ + __VA_ARGS__ \ + } + +#define REGISTER_KEY_TRANSFORM(params, k, ...) \ + template \ + __device__ __forceinline__ T KeyTransform(const Params& params, void* k_smem) { \ + __VA_ARGS__ \ + } + +#define REGISTER_LOGITS_TRANSFORM(params, logits, batch_idx, qo_idx, kv_idx, qo_head_idx, \ + kv_head_idx, ...) \ + template \ + __device__ __forceinline__ T LogitsTransform(const Params& params, T logits, uint32_t batch_idx, \ + uint32_t qo_idx, uint32_t kv_idx, \ + uint32_t qo_head_idx, uint32_t kv_head_idx) { \ + __VA_ARGS__ \ + } + +#define REGISTER_LOGITS_MASK(params, batch_idx, qo_idx, kv_idx, qo_head_idx, kv_head_idx, ...) \ + template \ + __device__ __forceinline__ bool LogitsMask(const Params& params, uint32_t batch_idx, \ + uint32_t qo_idx, uint32_t kv_idx, \ + uint32_t qo_head_idx, uint32_t kv_head_idx) { \ + __VA_ARGS__ \ + } + +#define REGISTER_M_D_UPDATE(params, kv_tile_idx, qo_head_idx, m, d, scale, ...) \ + template \ + __device__ __forceinline__ void update_m_d(const Params& params, uint32_t kv_tile_idx, \ + uint32_t qo_head_idx, T_M& m, float& d, \ + float& scale) { \ + __VA_ARGS__ \ + } + +#define REGISTER_OUTPUT_TRANSFORM(params, output, batch_idx, qo_idx, qo_head_idx, m, d, scale, \ + ...) \ + template \ + __device__ __forceinline__ T OutputTransform(const Params& params, T output, uint32_t batch_idx, \ + uint32_t qo_idx, uint32_t qo_head_idx, T_M& m, \ + float& d, float scale) { \ + __VA_ARGS__ \ + } + +struct AttentionVariantBase { + constexpr static bool use_softmax = true; + REGISTER_LOGITS_TRANSFORM(params, logits, batch_idx, qo_idx, kv_idx, qo_head_idx, kv_head_idx, + { return logits; }) + + REGISTER_LOGITS_MASK(params, batch_idx, qo_idx, kv_idx, qo_head_idx, kv_head_idx, + { return true; }) + + REGISTER_M_D_UPDATE(params, kv_tile_idx, qo_head_idx, m, d, scale, { return; }) + + REGISTER_OUTPUT_TRANSFORM(params, output, batch_idx, qo_idx, qo_head_idx, m, d, scale, { + float d_rcp = (m != -math::inf) ? math::ptx_rcp(d) : 0.f; + return output * d_rcp; + }) +}; + +} // namespace flashinfer + +#endif // FLASHINFER_ATTENTION_VARIANT_HELPER_H diff --git a/include/flashinfer/attention/variants.cuh b/include/flashinfer/attention/variants.cuh new file mode 100644 index 0000000000000000000000000000000000000000..e6e2a96707888bb9d6cd32cd71f9d60998e03f7d --- /dev/null +++ b/include/flashinfer/attention/variants.cuh @@ -0,0 +1,99 @@ +/* + * Copyright (c) 2024 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef FLASHINFER_ATTENTION_VARIANTS_CUH_ +#define FLASHINFER_ATTENTION_VARIANTS_CUH_ +#include + +#include +#include + +#include "../math.cuh" +#include "../utils.cuh" +#include "variant_helper.cuh" + +namespace flashinfer { + +DEFINE_HAS_MEMBER(maybe_mask_indptr) + +template +struct DefaultAttention : AttentionVariantBase { + static constexpr bool use_softmax = true; + + uint8_t* custom_mask_ptr; + uint32_t qo_len, kv_len; + uint32_t window_left; + float sm_scale_log2; + float soft_cap_pre_tanh_scale; + + // Create closure + template + __device__ __host__ DefaultAttention(const Params& params, uint32_t batch_idx, + uint8_t* smem_ptr) { + qo_len = params.get_qo_len(batch_idx); + kv_len = params.get_kv_len(batch_idx); + if constexpr (use_logits_soft_cap) { + soft_cap_pre_tanh_scale = params.sm_scale * math::ptx_rcp(params.logits_soft_cap); + sm_scale_log2 = math::log2e * params.logits_soft_cap; + } else { + if constexpr (use_alibi) { + sm_scale_log2 = math::log2e; + } else { + sm_scale_log2 = params.sm_scale * math::log2e; + } + } + if constexpr (use_custom_mask) { + if constexpr (has_maybe_mask_indptr_v) { + custom_mask_ptr = params.maybe_custom_mask + params.maybe_mask_indptr[batch_idx]; + } else { + custom_mask_ptr = params.maybe_custom_mask; + } + } + if constexpr (use_sliding_window) { + window_left = (params.window_left >= 0) ? params.window_left : kv_len; + } + } + + REGISTER_LOGITS_TRANSFORM(params, logits, batch_idx, qo_idx, kv_idx, qo_head_idx, kv_head_idx, { + if constexpr (use_alibi) { + logits = logits * params.sm_scale + + params.maybe_alibi_slopes[qo_head_idx] * float(int(kv_idx) - int(qo_idx)); + } + if constexpr (use_logits_soft_cap) { + logits = float(math::tanh(logits * soft_cap_pre_tanh_scale)); + } + return logits; + }) + + REGISTER_LOGITS_MASK(params, batch_idx, qo_idx, kv_idx, qo_head_idx, kv_head_idx, { + bool mask = true; + if constexpr (use_custom_mask) { + if (qo_idx >= qo_len || kv_idx >= kv_len) { + mask = false; + } else { + const uint64_t offset = static_cast(qo_idx) * kv_len + kv_idx; + mask &= ((custom_mask_ptr[offset / 8] >> (offset % 8)) & 1); + } + } + if constexpr (use_sliding_window) { + mask &= (kv_idx + qo_len + window_left >= kv_len + qo_idx); + } + return mask; + }) +}; + +}; // namespace flashinfer + +#endif // FLASHINFER_ATTENTION_VARIANTS_CUH_ diff --git a/include/flashinfer/attention_impl.cuh b/include/flashinfer/attention_impl.cuh new file mode 100644 index 0000000000000000000000000000000000000000..7ccfb74442001e49dd2d3e726b81902f18e4eae7 --- /dev/null +++ b/include/flashinfer/attention_impl.cuh @@ -0,0 +1,26 @@ +/* + * Copyright (c) 2024 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef FLASHINFER_ATTENTION_IMPL_CUH_ +#define FLASHINFER_ATTENTION_IMPL_CUH_ + +#include "attention/cascade.cuh" +#include "attention/decode.cuh" +#include "attention/default_decode_params.cuh" +#include "attention/default_prefill_params.cuh" +#include "attention/prefill.cuh" +#include "attention/variants.cuh" + +#endif // FLASHINFER_ATTENTION_IMPL_CUH_ diff --git a/include/flashinfer/comm/trtllm_allreduce.cuh b/include/flashinfer/comm/trtllm_allreduce.cuh new file mode 100644 index 0000000000000000000000000000000000000000..72566385040448636c60609d95151859067dc8a2 --- /dev/null +++ b/include/flashinfer/comm/trtllm_allreduce.cuh @@ -0,0 +1,1738 @@ +/* + * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include + +#include +#include + +#include "../exception.h" +#include "../logging.h" +#include "../utils.cuh" +#include "../vec_dtypes.cuh" + +namespace flashinfer { + +namespace trtllm_allreduce { + +constexpr size_t WARP_SIZE = 32; +constexpr size_t MAX_ALL_REDUCE_BLOCKS = 24; +constexpr size_t MAX_RANKS_PER_NODE = 16; +constexpr size_t DEFAULT_BLOCK_SIZE = 512; +constexpr size_t NUM_POINTERS_PER_RANK = 7; + +namespace details { + +static constexpr int kBytesPerAccess = 16; +static constexpr int kWarpSize = 32; +static constexpr int kMaxCtaSize = 1024; +static constexpr int kClusterMaxSize = 8; +static constexpr int kLamportTokenNumThreshold = 16; +static constexpr int kLamportHiddenSizeThreshold = 256; + +} // namespace details + +enum class AllReduceStrategyType : int8_t { + NCCL = 0, + MIN_LATENCY = 1, + UB = 2, + AUTO = 3, + ONESHOT = 4, + TWOSHOT = 5, + LOWPRECISION = 6, +}; + +enum class AllReduceStrategyConfig : int8_t { + USE_MEMCPY = 1 << 0, + PUSH_MODE = 1 << 1, +}; + +////////////////////// + +enum class AllReduceFusionOp : int8_t { + NONE = 0, + RESIDUAL_RMS_NORM = 1, + LAST_PROCESS_FOR_UB = 2, + RESIDUAL_RMS_PREPOST_NORM = 3, + RESIDUAL_RMS_NORM_QUANT_FP8 = 4, + RESIDUAL_RMS_NORM_QUANT_NVFP4 = 5, + RESIDUAL_RMS_NORM_OUT_QUANT_FP8 = 6, + RESIDUAL_RMS_NORM_OUT_QUANT_NVFP4 = 7, + MOE_ALLREDUCE_RESIDUAL_RMS_NORM = 8, + MOE_FINALIZE_ALLREDUCE_RESIDUAL_RMS_NORM = 9, +}; + +template +bool is_lamport_supported(int token_num, int hidden_size) { + if (!std::is_same_v && !std::is_same_v) { + return false; + } + if (token_num > details::kLamportTokenNumThreshold) { + return false; + } + if (hidden_size < details::kLamportHiddenSizeThreshold) { + return false; + } + return true; +} + +struct AllReduceFusionParams { + AllReduceFusionParams() + : bias_buffer(nullptr), + residual_buffer(nullptr), + weight_buffer(nullptr), + weight_buffer_pre_residual_norm(nullptr), + intermediate_buffer(nullptr) {} + + // gemm bias + void const* bias_buffer; + // residuial add + void const* residual_buffer; + // rms norm + int hidden_size; // equal to normalized_shape + void const* weight_buffer; // norm elem-wise affine gamma + void const* weight_buffer_pre_residual_norm; // for gemma norm before residual + float eps; + // new residual + void* intermediate_buffer; + void* lamport_peer_comm_buffer_ptrs[MAX_RANKS_PER_NODE * 3]; +}; + +template +struct AllReduceParams { + size_t elts_total; + size_t elts_per_rank; + size_t elts_per_block; + size_t rank_offset; + size_t ranks_per_node; + size_t local_rank; + uint32_t barrier_flag; + uint32_t* peer_barrier_ptrs_in[MAX_RANKS_PER_NODE]; + uint32_t* peer_barrier_ptrs_out[MAX_RANKS_PER_NODE]; + void* peer_comm_buffer_ptrs[MAX_RANKS_PER_NODE]; + void* local_output_buffer_ptr; + void const* local_input_buffer_ptr; + + AllReduceFusionParams fusion_params; + + static AllReduceParams deserialize(int64_t* buffer, size_t tpSize, size_t tpRank, int token_num, + int hidden_size, AllReduceFusionOp op) { + void* const* buffer_ptrs = reinterpret_cast(buffer); + int flag_offset; + if (op == AllReduceFusionOp::RESIDUAL_RMS_NORM && + is_lamport_supported(token_num, hidden_size)) { + flag_offset = 0; + } else { + flag_offset = 1; + } + auto const flag_ptr = &buffer[NUM_POINTERS_PER_RANK * tpSize + flag_offset]; + // cannot use 0 since 0 represents released state for barrier + *flag_ptr += 1; + uint32_t flag_value = *flag_ptr; + AllReduceParams params; + // Even plugins use ping buffers, odd plugins use pong. + // That way, we don't need to wait for other GPUs to be done + // before copying input tensor to workspace. + auto const buffer_offset = (flag_value % 2 == 0) ? 0 : tpSize; + + for (int i = 0; i < tpSize; ++i) { + params.peer_comm_buffer_ptrs[i] = buffer_ptrs[buffer_offset + i]; + } + for (int i = 0; i < tpSize; ++i) { + params.peer_barrier_ptrs_in[i] = reinterpret_cast(buffer_ptrs[2 * tpSize + i]); + } + for (int i = 0; i < tpSize; ++i) { + params.peer_barrier_ptrs_out[i] = reinterpret_cast(buffer_ptrs[3 * tpSize + i]); + } + params.barrier_flag = flag_value; + params.ranks_per_node = tpSize; + params.local_rank = tpRank; + + return params; + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct neg_zero { + static constexpr T value = -T(0); +}; + +template <> +struct neg_zero { + static constexpr unsigned short neg_zero_bits = 0x8000U; + static constexpr __half value = __half_raw{neg_zero_bits}; +}; + +template <> +struct neg_zero { + static constexpr unsigned short neg_zero_bits = 0x8000U; + static constexpr __nv_bfloat16 value = __nv_bfloat16_raw{neg_zero_bits}; +}; + +template <> +struct neg_zero { + static constexpr unsigned int neg_zero_bits = 0x80000000U; + static constexpr float value = -0.0f; +}; + +template +__device__ static constexpr T neg_zero_v = neg_zero::value; + +template +__device__ bool is_negative_zero(T) { + return false; +} + +// float specialization +template <> +__device__ bool is_negative_zero(float x) { + return (__float_as_int(x) == 0x80000000); +} + +// double specialization +template <> +__device__ bool is_negative_zero(double x) { + return (__double_as_longlong(x) == 0x8000000000000000ULL); +} + +// __half specialization +template <> +__device__ bool is_negative_zero<__half>(__half x) { + return (__half_as_ushort(x) == 0x8000); +} + +// __nv_bfloat16 specialization +template <> +__device__ bool is_negative_zero<__nv_bfloat16>(__nv_bfloat16 x) { + return (__bfloat16_as_ushort(x) == 0x8000); +} + +template +__device__ __forceinline__ bool has_neg_zero(const vec_t& vec) { +#pragma unroll + for (int i = 0; i < VEC_SIZE; ++i) { + if (is_negative_zero(vec[i])) { + return true; + } + } + return false; +} + +template +__device__ __forceinline__ void remove_neg_zero(vec_t& vec) { +#pragma unroll + for (int i = 0; i < VEC_SIZE; ++i) { + vec[i] = (is_negative_zero(vec[i])) ? static_cast(0.f) : vec[i]; + } +} + +template +__device__ __forceinline__ void set_neg_zero(T* addr) { + vec_t val; + val.fill(neg_zero_v); + val.store_global_volatile(addr); +} + +static inline __device__ void st_flag_release(uint32_t const& flag, uint32_t* flag_addr) { +#if __CUDA_ARCH__ >= 700 + asm volatile("st.global.release.sys.b32 [%1], %0;" ::"r"(flag), "l"(flag_addr)); +#else + __threadfence_system(); + asm volatile("st.global.volatile.b32 [%1], %0;" ::"r"(flag), "l"(flag_addr)); +#endif +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +static inline __device__ uint32_t ld_flag_acquire(uint32_t* flag_addr) { + uint32_t flag; +#if __CUDA_ARCH__ >= 700 + asm volatile("ld.global.acquire.sys.b32 %0, [%1];" : "=r"(flag) : "l"(flag_addr)); +#else + asm volatile("ld.global.volatile.b32 %0, [%1];" : "=r"(flag) : "l"(flag_addr)); +#endif + return flag; +} + +template +__device__ __forceinline__ vec_t vec_add(const vec_t& a, + const vec_t& b) { + vec_t ret; +#pragma unroll + for (int i = 0; i < VEC_SIZE; ++i) { + ret[i] = static_cast(a[i]) + static_cast(b[i]); + } + return ret; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +__inline__ __device__ void multi_gpu_barrier(uint32_t** signals, uint32_t const flag, + size_t const local_rank, size_t const world_size, + int const tidx, int const bidx) { + // After this function, at least one block in each GPU has reached the barrier + if (tidx < world_size) { + // we can think of signals having the shape [world_size, world_size] + // Dimension 0 is the "listening" dimension, dimension 1 is "emitting" dimension + + // Block 0 broadcasts its flag (local_rank on emitting dimension) to all receivers + size_t offset = (flag % 2) ? world_size : 0; + + if (bidx == 0) { + st_flag_release(flag, signals[tidx] + offset + local_rank); + } + + // All blocks check that corresponding block 0 on other GPUs have set the flag + // No deadlock because block #0 is always the first block started + uint32_t* peer_barrier_d = signals[local_rank] + offset + tidx; + while (ld_flag_acquire(peer_barrier_d) != flag) { + } + } + + __syncthreads(); +} + +__inline__ __device__ void block_barrier(uint32_t** signals, uint32_t const flag, + size_t const local_rank, size_t const world_size, + int const tidx, int const bidx, int const grid_size) { + // After this function, the block of id == bidx of each GPU has reached the barrier + if (tidx < world_size) { + // we can think of signals having the shape [world_size, 2, num_blocks, world_size] + // (+ an offset on dim 2 to account for flags used in multi_gpu_barrier) + // Dimension 0 is the "listening" dimension, dimension 3 is "emitting" dimension + + // Block broadcast its flag (local_rank on emitting dimension) to all receivers + uint32_t flag_block_offset = world_size + bidx * world_size; + + if (flag % 2 == 1) { + flag_block_offset += (grid_size + 1) * world_size; + } + + st_flag_release(flag, signals[tidx] + flag_block_offset + local_rank); + + // Blocks check that corresponding blocks on other GPUs have also set the flag + uint32_t* peer_barrier_d = signals[local_rank] + flag_block_offset + tidx; + + while (ld_flag_acquire(peer_barrier_d) != flag) { + } + } + + __syncthreads(); +} + +namespace reduce_fusion { + +inline __device__ float warp_reduce_sum(float val) { + val += __shfl_xor_sync(~0, val, 16); + val += __shfl_xor_sync(~0, val, 8); + val += __shfl_xor_sync(~0, val, 4); + val += __shfl_xor_sync(~0, val, 2); + val += __shfl_xor_sync(~0, val, 1); + return val; +} + +inline __device__ float block_reduce_sum(float val) { + __shared__ float smem[details::kWarpSize]; + int lane_id = threadIdx.x % details::kWarpSize, warp_id = threadIdx.x / details::kWarpSize, + warp_num = blockDim.x / details::kWarpSize; + val = warp_reduce_sum(val); + if (lane_id == 0) { + smem[warp_id] = val; + } + __syncthreads(); + val = lane_id < warp_num ? smem[lane_id] : 0.f; + val = warp_reduce_sum(val); + return val; +} + +template +inline __device__ float accumulate(float acc, vec_t& vec) { +#pragma unroll + for (int i = 0; i < VEC_SIZE; ++i) { + float v = static_cast(vec[i]); + acc += v * v; + } + return acc; +} + +template +inline __device__ vec_t rms_norm(float denom, vec_t& vec, + vec_t& weight) { + vec_t ret; +#pragma unroll + for (int i = 0; i < VEC_SIZE; ++i) { + float v1 = static_cast(vec[i]); + if constexpr (Affine) { + float v2 = static_cast(weight[i]); + ret[i] = static_cast(v1 * denom * v2); + } else { + ret[i] = static_cast(v1 * denom); + } + } + return ret; +} + +template +__global__ void rms_norm_kernel(AllReduceParams params) { + static constexpr uint32_t VEC_SIZE = 16 / sizeof(T); + + extern __shared__ uint8_t smem_ptr[]; + T* smem = reinterpret_cast(smem_ptr); + + int bid = blockIdx.x, tid = threadIdx.x; + + T const* bias_buffer = reinterpret_cast(params.fusion_params.bias_buffer); + T const* residual_buffer = reinterpret_cast(params.fusion_params.residual_buffer); + T const* weight_buffer = reinterpret_cast(params.fusion_params.weight_buffer); + T* local_final_output_buffer = reinterpret_cast(params.local_output_buffer_ptr); + T* intermediate_buffer = reinterpret_cast(params.fusion_params.intermediate_buffer); + + int block_offset = bid * params.fusion_params.hidden_size; + int thread_offset = tid * VEC_SIZE; + + if constexpr (Residual) { + residual_buffer += block_offset; + } + local_final_output_buffer += block_offset; + intermediate_buffer += block_offset; + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) && (__CUDA_ARCH__ < 1200)) + cudaGridDependencySynchronize(); +#endif + + vec_t inter_vec, weight_vec; + float acc = 0.f; + for (int offset = thread_offset; offset < params.fusion_params.hidden_size; + offset += blockDim.x * VEC_SIZE) { + inter_vec.load(intermediate_buffer + offset); + if constexpr (Bias) { + vec_t bias_vec; + bias_vec.load(bias_buffer + offset); + inter_vec = vec_add(inter_vec, bias_vec); + } + if constexpr (Residual) { + vec_t residual_vec; + residual_vec.load(residual_buffer + offset); + inter_vec = vec_add(inter_vec, residual_vec); + inter_vec.store(intermediate_buffer + offset); + } + acc = accumulate(acc, inter_vec); + if constexpr (UseSmem) { + inter_vec.store(&smem[offset]); + } + } + acc = block_reduce_sum(acc); + float denom = rsqrtf(acc / params.fusion_params.hidden_size + params.fusion_params.eps); + for (int offset = thread_offset; offset < params.fusion_params.hidden_size; + offset += blockDim.x * VEC_SIZE) { + if constexpr (UseSmem) { + inter_vec.load(&smem[offset]); + } + if constexpr (Affine) { + weight_vec.load(weight_buffer + offset); + } + inter_vec = rms_norm(denom, inter_vec, weight_vec); + inter_vec.store(&local_final_output_buffer[offset]); + } +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) && (__CUDA_ARCH__ < 1200)) + cudaTriggerProgrammaticLaunchCompletion(); +#endif +} + +template +__global__ void rms_pre_post_norm_kernel( + AllReduceParams params) // for gemma2 pre residual + post residual norm +{ + static constexpr uint32_t VEC_SIZE = 16 / sizeof(T); + + int bid = blockIdx.x, tid = threadIdx.x; + + T const* bias_buffer = reinterpret_cast(params.fusion_params.bias_buffer); + T const* residual_buffer = reinterpret_cast(params.fusion_params.residual_buffer); + T const* weight_buffer = reinterpret_cast(params.fusion_params.weight_buffer); + T const* weight_buffer_pre_residual_norm = + reinterpret_cast(params.fusion_params.weight_buffer_pre_residual_norm); + T* local_final_output_buffer = reinterpret_cast(params.local_output_buffer_ptr); + T* intermediate_buffer = reinterpret_cast(params.fusion_params.intermediate_buffer); + + int block_offset = bid * params.fusion_params.hidden_size; + int thread_offset = tid * VEC_SIZE; + + if constexpr (Residual) { + residual_buffer += block_offset; + } + local_final_output_buffer += block_offset; + intermediate_buffer += block_offset; + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) && (__CUDA_ARCH__ < 1200)) + cudaGridDependencySynchronize(); +#endif + + vec_t inter_vec, weight_vec, weight_vec_pre_residual_norm, bias_vec; + float acc = 0.f; + float acc_pre_residual_norm = 0.f; + for (int offset = thread_offset; offset < params.fusion_params.hidden_size; + offset += blockDim.x * VEC_SIZE) { + inter_vec.load(intermediate_buffer + offset); + if constexpr (Bias) { + bias_vec.load(bias_buffer + offset); + } + + if constexpr (Bias) { + inter_vec = vec_add(inter_vec, bias_vec); + } + + // pre-residual norm. + acc_pre_residual_norm = accumulate(acc_pre_residual_norm, inter_vec); + acc_pre_residual_norm = block_reduce_sum(acc_pre_residual_norm); + float denom_pre_residual_norm = + rsqrtf(acc_pre_residual_norm / params.fusion_params.hidden_size + params.fusion_params.eps); + + if constexpr (Affine) { + weight_vec_pre_residual_norm.load(weight_buffer_pre_residual_norm + thread_offset); + } + inter_vec = rms_norm(denom_pre_residual_norm, inter_vec, + weight_vec_pre_residual_norm); + + if constexpr (Residual) { + vec_t residual_vec; + residual_vec.load(residual_buffer + offset); + inter_vec = vec_add(inter_vec, residual_vec); + inter_vec.store(intermediate_buffer + offset); + } + acc = accumulate(acc, inter_vec); + } + acc = block_reduce_sum(acc); + float denom = rsqrtf(acc / params.fusion_params.hidden_size + params.fusion_params.eps); + for (int offset = thread_offset; offset < params.fusion_params.hidden_size; + offset += blockDim.x * VEC_SIZE) { + if constexpr (Affine) { + weight_vec.load(weight_buffer + offset); + } + inter_vec = rms_norm(denom, inter_vec, weight_vec); + inter_vec.store(&local_final_output_buffer[offset]); + } +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) && (__CUDA_ARCH__ < 1200)) + cudaTriggerProgrammaticLaunchCompletion(); +#endif +} + +template +cudaError_t rms_norm_kernel_launcher(AllReduceParams& params, AllReduceFusionOp fusionOp, + bool launch_with_pdl, cudaStream_t stream) { + static constexpr uint32_t VEC_SIZE = 16 / sizeof(T); + FLASHINFER_CHECK(params.fusion_params.hidden_size % VEC_SIZE == 0, + "hidden_size must be a multiple of ", VEC_SIZE); + if (fusionOp == AllReduceFusionOp::RESIDUAL_RMS_PREPOST_NORM) { + FLASHINFER_CHECK(params.fusion_params.hidden_size <= 8192, + "hidden_size must be less than or equal to 8192"); + } + int need_threads = params.fusion_params.hidden_size / VEC_SIZE; + int cta_size; + if (need_threads <= details::kMaxCtaSize) { + cta_size = (need_threads + details::kWarpSize - 1) / details::kWarpSize * details::kWarpSize; + } else { + cta_size = details::kMaxCtaSize; + } + int cta_num = params.elts_total / params.fusion_params.hidden_size; + int smem_size = 0; + if (cta_size * details::kBytesPerAccess / sizeof(T) < params.fusion_params.hidden_size) { + smem_size = params.fusion_params.hidden_size * sizeof(T); + cudaLaunchConfig_t kernelConfig = {0}; + kernelConfig.gridDim = cta_num; + kernelConfig.blockDim = cta_size; + kernelConfig.dynamicSmemBytes = smem_size; + kernelConfig.stream = stream; + + cudaLaunchAttribute attribute[1]; + attribute[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attribute[0].val.programmaticStreamSerializationAllowed = launch_with_pdl; + kernelConfig.attrs = attribute; + kernelConfig.numAttrs = 1; + + if (fusionOp == AllReduceFusionOp::RESIDUAL_RMS_NORM) { + FLASHINFER_CUDA_CALL(cudaLaunchKernelEx( + &kernelConfig, rms_norm_kernel, params)); + } else { // AllReduceFusionOp::RESIDUAL_RMS_PREPOST_NORM + FLASHINFER_CUDA_CALL(cudaLaunchKernelEx( + &kernelConfig, rms_pre_post_norm_kernel, params)); + } + } else { + cudaLaunchConfig_t kernelConfig = {0}; + kernelConfig.gridDim = cta_num; + kernelConfig.blockDim = cta_size; + kernelConfig.dynamicSmemBytes = smem_size; + kernelConfig.stream = stream; + + cudaLaunchAttribute attribute[1]; + attribute[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attribute[0].val.programmaticStreamSerializationAllowed = launch_with_pdl; + kernelConfig.attrs = attribute; + kernelConfig.numAttrs = 1; + + if (fusionOp == AllReduceFusionOp::RESIDUAL_RMS_NORM) { + FLASHINFER_CUDA_CALL(cudaLaunchKernelEx( + &kernelConfig, rms_norm_kernel, params)); + } else { // AllReduceFusionOp::RESIDUAL_RMS_PREPOST_NORM + FLASHINFER_CUDA_CALL(cudaLaunchKernelEx( + &kernelConfig, rms_pre_post_norm_kernel, params)); + } + } + return cudaSuccess; +} + +template +struct Reducer; + +template +struct Reducer { + constexpr static uint32_t VEC_SIZE = 16 / sizeof(T); + static __device__ __forceinline__ vec_t allreduce(AllReduceParams& params, + int global_offset) { + int ping = params.barrier_flag % 3; + int pong = (params.barrier_flag + 2) % 3; + T const* local_input_buffer = reinterpret_cast(params.local_input_buffer_ptr); + T* local_shared_buffer = reinterpret_cast( + params.fusion_params + .lamport_peer_comm_buffer_ptrs[params.local_rank + ping * MAX_RANKS_PER_NODE]); + T* local_clean_buffer = reinterpret_cast( + params.fusion_params + .lamport_peer_comm_buffer_ptrs[params.local_rank + pong * MAX_RANKS_PER_NODE]); + local_input_buffer += global_offset; + local_shared_buffer += global_offset; + local_clean_buffer += global_offset; + T* buffers[RanksPerNode]; +#pragma unroll + for (int ii = 0; ii < RanksPerNode; ++ii) { + int rank = (params.local_rank + ii) % RanksPerNode; + buffers[ii] = reinterpret_cast( + params.fusion_params + .lamport_peer_comm_buffer_ptrs[rank + ping * MAX_RANKS_PER_NODE]) + + global_offset + params.local_rank * params.elts_total; + } + vec_t sum_vec, val; + val.load(local_input_buffer); +#pragma unroll + for (int ii = 1; ii < RanksPerNode; ++ii) { + val.store_global_volatile(buffers[ii]); + } + sum_vec = val; +#pragma unroll + for (int ii = 1; ii < RanksPerNode; ++ii) { + int rank = (params.local_rank + ii) % RanksPerNode; + set_neg_zero(local_clean_buffer + rank * params.elts_total); + } + vec_t vals[RanksPerNode - 1]; + bool done = false; + while (!done) { + done = true; +#pragma unroll + for (int ii = 1; ii < RanksPerNode; ++ii) { + int rank = (params.local_rank + ii) % RanksPerNode; + vals[ii - 1].load_global_volatile(local_shared_buffer + rank * params.elts_total); + } +#pragma unroll + for (int ii = 0; ii < RanksPerNode - 1; ii++) { + done &= !has_neg_zero(vals[ii]); + } + } + +#pragma unroll + for (int ii = 1; ii < RanksPerNode; ++ii) { + sum_vec = vec_add(sum_vec, vals[ii - 1]); + } + return sum_vec; + } +}; + +template +struct Reducer { + static constexpr uint32_t VEC_SIZE = 16 / sizeof(T); + static __device__ __forceinline__ vec_t allreduce(AllReduceParams& params, + int global_offset) { + int ping = params.barrier_flag % 3; + int pong = (params.barrier_flag + 2) % 3; + T const* local_input_buffer = reinterpret_cast(params.local_input_buffer_ptr); + T* local_shared_buffer = reinterpret_cast( + params.fusion_params + .lamport_peer_comm_buffer_ptrs[params.local_rank + ping * MAX_RANKS_PER_NODE]); + T* local_clean_buffer = reinterpret_cast( + params.fusion_params + .lamport_peer_comm_buffer_ptrs[params.local_rank + pong * MAX_RANKS_PER_NODE]); + local_input_buffer += global_offset; + local_shared_buffer += global_offset; + local_clean_buffer += global_offset; + T* buffers[RanksPerNode]; +#pragma unroll + for (int ii = 0; ii < RanksPerNode; ++ii) { + int rank = (params.local_rank + ii) % RanksPerNode; + buffers[ii] = reinterpret_cast( + params.fusion_params + .lamport_peer_comm_buffer_ptrs[rank + ping * MAX_RANKS_PER_NODE]) + + global_offset; + } + vec_t sum_vec, val; + val.load(local_input_buffer); + val.store_global_volatile(reinterpret_cast(local_shared_buffer)); + sum_vec = val; +#pragma unroll + for (int ii = 1; ii < RanksPerNode; ++ii) { + do { + val.load_global_volatile(reinterpret_cast(buffers[ii])); + } while (has_neg_zero(val)); + sum_vec = vec_add(sum_vec, val); + } + set_neg_zero(local_clean_buffer); + return sum_vec; + } +}; + +template +__global__ void lamport_style_one_shot_all_reduce_norm_kernel(AllReduceParams params) { +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) && (__CUDA_ARCH__ < 1200)) + namespace cg = cooperative_groups; + static_assert(RanksPerNode <= MAX_RANKS_PER_NODE); + static constexpr uint32_t VEC_SIZE = 16 / sizeof(T); + + cg::cluster_group cluster = cg::this_cluster(); + + __shared__ float cluster_acc, cluster_acc_sum; + + int bid = blockIdx.x, tid = threadIdx.x; + int cluster_id = bid / ClusterSize, cluster_block_rank = bid % ClusterSize; + + int token_id = cluster_id; + int cluster_offset = token_id * params.fusion_params.hidden_size; + int block_offset = cluster_block_rank * params.fusion_params.hidden_size / ClusterSize; + int thread_offset = tid * VEC_SIZE; + + int inner_token_offset = block_offset + thread_offset; + int global_offset = cluster_offset + inner_token_offset; + + T const* bias_buffer = reinterpret_cast(params.fusion_params.bias_buffer); + T const* residual_buffer = reinterpret_cast(params.fusion_params.residual_buffer); + T const* weight_buffer = reinterpret_cast(params.fusion_params.weight_buffer); + T* local_final_output_buffer = reinterpret_cast(params.local_output_buffer_ptr); + T* intermediate_buffer = reinterpret_cast(params.fusion_params.intermediate_buffer); + + local_final_output_buffer += global_offset; + intermediate_buffer += global_offset; + residual_buffer += global_offset; + bias_buffer += inner_token_offset; + weight_buffer += inner_token_offset; + + vec_t weight_vec, bias_vec, residual_vec; + residual_vec.load(residual_buffer); + if constexpr (Bias) { + bias_vec.load(bias_buffer); + } + if constexpr (Affine) { + weight_vec.load(weight_buffer); + } + + cudaGridDependencySynchronize(); + + float acc = 0.f; + vec_t sum_vec; + sum_vec = Reducer::allreduce(params, global_offset); + + if constexpr (Bias) { + sum_vec = vec_add(sum_vec, bias_vec); + } + sum_vec = vec_add(sum_vec, residual_vec); + sum_vec.store(intermediate_buffer); + acc = accumulate(acc, sum_vec); + acc = block_reduce_sum(acc); + if (ClusterSize > 1) { + if (threadIdx.x == 0) { + cluster_acc = acc; + } + cluster.sync(); + if (threadIdx.x == 0) { + acc = 0.f; +#pragma unroll + for (int ii = 0; ii < ClusterSize; ++ii) { + acc += *cluster.map_shared_rank(&cluster_acc, ii); + } + cluster_acc_sum = acc; + } + __syncthreads(); + acc = cluster_acc_sum; + cluster.sync(); + } + + float denom = rsqrtf(acc / params.fusion_params.hidden_size + params.fusion_params.eps); + sum_vec = rms_norm(denom, sum_vec, weight_vec); + sum_vec.store(local_final_output_buffer); + + cudaTriggerProgrammaticLaunchCompletion(); +#endif +} + +int heuristic_min_warp_number(int tp_size, int hidden_size) { + if (hidden_size >= 4096) { + return 4; + } + if (tp_size == 2) { + return 32; + } else { + return 16; + } +} + +template +cudaError_t lamport_style_one_shot_all_reduce_norm_kernel_launcher(AllReduceParams params, + bool launch_with_pdl, + cudaStream_t stream) { + static constexpr uint32_t VEC_SIZE = 16 / sizeof(T); + FLASHINFER_CHECK(params.fusion_params.hidden_size % VEC_SIZE == 0, + "hidden_size must be a multiple of ", VEC_SIZE); + int threads_per_token = params.fusion_params.hidden_size / VEC_SIZE; + int warps_per_token = (threads_per_token + details::kWarpSize - 1) / details::kWarpSize; + int token_num = params.elts_total / params.fusion_params.hidden_size; + int warp_min_number = heuristic_min_warp_number(RanksPerNode, params.fusion_params.hidden_size); + int cluster_size = std::min(((warps_per_token + warp_min_number - 1) / warp_min_number), + details::kClusterMaxSize); + int cta_size = warps_per_token / cluster_size * details::kWarpSize; + FLASHINFER_CHECK(cta_size <= details::kMaxCtaSize, "cta_size must be less than or equal to ", + details::kMaxCtaSize); + int cta_num = token_num * cluster_size; + cudaLaunchConfig_t kernel_config = {0}; + kernel_config.gridDim = cta_num; + kernel_config.blockDim = cta_size; + kernel_config.dynamicSmemBytes = 0; + kernel_config.stream = stream; + + cudaLaunchAttribute attribute[2]; + attribute[0].id = cudaLaunchAttributeClusterDimension; + attribute[0].val.clusterDim.x = cluster_size; + attribute[0].val.clusterDim.y = 1; + attribute[0].val.clusterDim.z = 1; + kernel_config.attrs = attribute; + kernel_config.numAttrs = 1; + if (launch_with_pdl) { + attribute[1].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attribute[1].val.programmaticStreamSerializationAllowed = 1; + kernel_config.numAttrs++; + } +#define LAUNCH_LAMPORT_KERNEL(CLUSTER_SIZE) \ + if (cluster_size == CLUSTER_SIZE) { \ + FLASHINFER_CUDA_CALL(cudaLaunchKernelEx( \ + &kernel_config, \ + lamport_style_one_shot_all_reduce_norm_kernel, \ + params)); \ + return cudaSuccess; \ + } + LAUNCH_LAMPORT_KERNEL(1); + LAUNCH_LAMPORT_KERNEL(2); + LAUNCH_LAMPORT_KERNEL(3); + LAUNCH_LAMPORT_KERNEL(4); + LAUNCH_LAMPORT_KERNEL(5); + LAUNCH_LAMPORT_KERNEL(6); + LAUNCH_LAMPORT_KERNEL(7); + LAUNCH_LAMPORT_KERNEL(8); +#undef LAUNCH_LAMPORT_KERNEL +} + +template +__global__ void __launch_bounds__(1024, 1) + one_shot_all_reduce_norm_kernel(AllReduceParams params) { + static constexpr uint32_t VEC_SIZE = 16 / sizeof(T); + + extern __shared__ uint8_t smem_ptr[]; + T* smem = reinterpret_cast(smem_ptr); + + int bid = blockIdx.x, tid = threadIdx.x; + int norm_num = params.elts_total / params.fusion_params.hidden_size; + int norm_per_block = (norm_num + gridDim.x - 1) / gridDim.x; + int norm_this_block = std::min(norm_per_block, norm_num - bid * norm_per_block); + + T const* local_input_buffer = reinterpret_cast(params.local_input_buffer_ptr); + T const* bias_buffer = reinterpret_cast(params.fusion_params.bias_buffer); + T const* residual_buffer = reinterpret_cast(params.fusion_params.residual_buffer); + T const* weight_buffer = reinterpret_cast(params.fusion_params.weight_buffer); + T* local_shared_buffer = reinterpret_cast(params.peer_comm_buffer_ptrs[params.local_rank]); + T* local_final_output_buffer = reinterpret_cast(params.local_output_buffer_ptr); + T* intermediate_buffer = reinterpret_cast(params.fusion_params.intermediate_buffer); + + int block_offset = bid * norm_per_block * params.fusion_params.hidden_size; + int thread_offset = tid * VEC_SIZE; + + local_input_buffer += block_offset; + residual_buffer += block_offset; + local_shared_buffer += block_offset; + local_final_output_buffer += block_offset; + intermediate_buffer += block_offset; + + T* buffers[RanksPerNode]; +#pragma unroll + for (int ii = 0; ii < RanksPerNode; ++ii) { + int rank = (params.local_rank + ii) % RanksPerNode; + buffers[ii] = reinterpret_cast(params.peer_comm_buffer_ptrs[rank]); + } + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) && (__CUDA_ARCH__ < 1200)) + cudaGridDependencySynchronize(); +#endif + + for (int offset = thread_offset; offset < norm_this_block * params.fusion_params.hidden_size; + offset += blockDim.x * VEC_SIZE) { + *reinterpret_cast(&local_shared_buffer[offset]) = + *reinterpret_cast(&local_input_buffer[offset]); + } + block_barrier(params.peer_barrier_ptrs_in, params.barrier_flag, params.local_rank, RanksPerNode, + tid, bid, gridDim.x); + for (int norm_idx = 0; norm_idx < norm_this_block; ++norm_idx) { + int norm_offset = norm_idx * params.fusion_params.hidden_size; + float acc = 0.f; + vec_t sum_vec, weight_vec, bias_vec, residual_vec; + for (int offset = thread_offset; offset < params.fusion_params.hidden_size; + offset += blockDim.x * VEC_SIZE) { + vec_t vals[RanksPerNode]; + sum_vec.fill(T(0)); + if constexpr (Bias) { + bias_vec.load(bias_buffer + offset); + } + residual_vec.load(residual_buffer + norm_offset + offset); +#pragma unroll + for (int ii = 0; ii < RanksPerNode; ++ii) { + vals[ii].load(buffers[ii] + block_offset + norm_offset + offset); + } +#pragma unroll + for (int ii = 0; ii < RanksPerNode; ++ii) { + sum_vec = vec_add(sum_vec, vals[ii]); + } + if constexpr (Bias) { + sum_vec = vec_add(sum_vec, bias_vec); + } + sum_vec = vec_add(sum_vec, residual_vec); + sum_vec.store(&intermediate_buffer[norm_offset + offset]); + acc = accumulate(acc, sum_vec); + if constexpr (UseSmem) { + sum_vec.store(&smem[offset]); + } + } + acc = block_reduce_sum(acc); + float denom = rsqrtf(acc / params.fusion_params.hidden_size + params.fusion_params.eps); + for (int offset = thread_offset; offset < params.fusion_params.hidden_size; + offset += blockDim.x * VEC_SIZE) { + if constexpr (UseSmem) { + sum_vec.load(&smem[offset]); + } + if constexpr (Affine) { + weight_vec.load(weight_buffer + offset); + } + sum_vec = rms_norm(denom, sum_vec, weight_vec); + sum_vec.store(&local_final_output_buffer[norm_offset + offset]); + } + } +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) && (__CUDA_ARCH__ < 1200)) + cudaTriggerProgrammaticLaunchCompletion(); +#endif +} + +template +__global__ void __launch_bounds__(1024, 1) + one_shot_prenorm_all_reduce_norm_kernel(AllReduceParams params) { + static constexpr uint32_t VEC_SIZE = 16 / sizeof(T); + + int bid = blockIdx.x, tid = threadIdx.x; + int norm_num = params.elts_total / params.fusion_params.hidden_size; + int norm_per_block = (norm_num + gridDim.x - 1) / gridDim.x; + int norm_this_block = std::min(norm_per_block, norm_num - bid * norm_per_block); + + T const* local_input_buffer = reinterpret_cast(params.local_input_buffer_ptr); + T const* bias_buffer = reinterpret_cast(params.fusion_params.bias_buffer); + T const* residual_buffer = reinterpret_cast(params.fusion_params.residual_buffer); + T const* weight_buffer = reinterpret_cast(params.fusion_params.weight_buffer); + T const* weight_buffer_pre_residual_norm = + reinterpret_cast(params.fusion_params.weight_buffer_pre_residual_norm); + T* local_shared_buffer = reinterpret_cast(params.peer_comm_buffer_ptrs[params.local_rank]); + T* local_final_output_buffer = reinterpret_cast(params.local_output_buffer_ptr); + T* intermediate_buffer = reinterpret_cast(params.fusion_params.intermediate_buffer); + + int block_offset = bid * norm_per_block * params.fusion_params.hidden_size; + int thread_offset = tid * VEC_SIZE; + + local_input_buffer += block_offset; + residual_buffer += block_offset; + local_shared_buffer += block_offset; + local_final_output_buffer += block_offset; + intermediate_buffer += block_offset; + + T* buffers[RanksPerNode]; +#pragma unroll + for (int ii = 0; ii < RanksPerNode; ++ii) { + int rank = (params.local_rank + ii) % RanksPerNode; + buffers[ii] = reinterpret_cast(params.peer_comm_buffer_ptrs[rank]); + } + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) && (__CUDA_ARCH__ < 1200)) + cudaGridDependencySynchronize(); +#endif + + for (int offset = thread_offset; offset < norm_this_block * params.fusion_params.hidden_size; + offset += blockDim.x * VEC_SIZE) { + *reinterpret_cast(&local_shared_buffer[offset]) = + *reinterpret_cast(&local_input_buffer[offset]); + } + block_barrier(params.peer_barrier_ptrs_in, params.barrier_flag, params.local_rank, RanksPerNode, + tid, bid, gridDim.x); + for (int norm_idx = 0; norm_idx < norm_this_block; ++norm_idx) { + int norm_offset = norm_idx * params.fusion_params.hidden_size; + float acc = 0.f; + float acc_pre_residual_norm = 0.f; + vec_t sum_vec, weight_vec, bias_vec, residual_vec, weight_vec_pre_residual_norm; + for (int offset = thread_offset; offset < params.fusion_params.hidden_size; + offset += blockDim.x * VEC_SIZE) { + vec_t vals[RanksPerNode]; + sum_vec.fill(T(0)); + if constexpr (Bias) { + bias_vec.load(bias_buffer + offset); + } + residual_vec.load(residual_buffer + norm_offset + offset); +#pragma unroll + for (int ii = 0; ii < RanksPerNode; ++ii) { + vals[ii].load(buffers[ii] + block_offset + norm_offset + offset); + } +#pragma unroll + for (int ii = 0; ii < RanksPerNode; ++ii) { + sum_vec = vec_add(sum_vec, vals[ii]); + } + + if constexpr (Bias) { + sum_vec = vec_add(sum_vec, bias_vec); + } + + // norm1 is pre-residual norm. + acc_pre_residual_norm = accumulate(acc_pre_residual_norm, sum_vec); + + acc_pre_residual_norm = block_reduce_sum(acc_pre_residual_norm); + + float denom_pre_residual_norm = rsqrtf( + acc_pre_residual_norm / params.fusion_params.hidden_size + params.fusion_params.eps); + if constexpr (Affine) { + weight_vec_pre_residual_norm.load(weight_buffer_pre_residual_norm + thread_offset); + } + sum_vec = rms_norm(denom_pre_residual_norm, sum_vec, + weight_vec_pre_residual_norm); + + sum_vec = vec_add(sum_vec, residual_vec); + sum_vec.store(&intermediate_buffer[norm_offset + offset]); + acc = accumulate(acc, sum_vec); + } + acc = block_reduce_sum(acc); + float denom = rsqrtf(acc / params.fusion_params.hidden_size + params.fusion_params.eps); + if constexpr (Affine) { + weight_vec.load(weight_buffer + thread_offset); + } + sum_vec = rms_norm(denom, sum_vec, weight_vec); + sum_vec.store(&local_final_output_buffer[norm_offset + thread_offset]); + } +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) && (__CUDA_ARCH__ < 1200)) + cudaTriggerProgrammaticLaunchCompletion(); +#endif +} + +template +cudaError_t one_shot_all_reduce_norm_kernel_launcher(AllReduceParams& params, + AllReduceFusionOp fusionOp, + bool launch_with_pdl, cudaStream_t stream) { + int token_num = params.elts_total / params.fusion_params.hidden_size; + + if (fusionOp == AllReduceFusionOp::RESIDUAL_RMS_PREPOST_NORM) { + FLASHINFER_CHECK(params.fusion_params.hidden_size <= 8192, + "hidden_size must be less than or equal to 8192"); + } + + if (is_lamport_supported(token_num, params.fusion_params.hidden_size) && + (fusionOp != AllReduceFusionOp::RESIDUAL_RMS_PREPOST_NORM)) { + lamport_style_one_shot_all_reduce_norm_kernel_launcher( + params, launch_with_pdl, stream); + } else { + static constexpr uint32_t VEC_SIZE = 16 / sizeof(T); + FLASHINFER_CHECK(params.fusion_params.hidden_size % VEC_SIZE == 0, + "hidden_size must be a multiple of ", VEC_SIZE); + int need_threads = params.fusion_params.hidden_size / VEC_SIZE; + int cta_size; + if (need_threads <= details::kMaxCtaSize) { + cta_size = (need_threads + details::kWarpSize - 1) / details::kWarpSize * details::kWarpSize; + } else { + cta_size = details::kMaxCtaSize; + } + int norm_num = params.elts_total / params.fusion_params.hidden_size; + int cta_num = std::min(norm_num, static_cast(MAX_ALL_REDUCE_BLOCKS)); + int smem_size = 0; + + if (cta_size * VEC_SIZE < params.fusion_params.hidden_size) { + smem_size = params.fusion_params.hidden_size * sizeof(T); + cudaLaunchConfig_t kernelConfig = {0}; + kernelConfig.gridDim = cta_num; + kernelConfig.blockDim = cta_size; + kernelConfig.dynamicSmemBytes = smem_size; + kernelConfig.stream = stream; + + cudaLaunchAttribute attribute[1]; + attribute[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attribute[0].val.programmaticStreamSerializationAllowed = launch_with_pdl; + kernelConfig.attrs = attribute; + kernelConfig.numAttrs = 1; + if (fusionOp == AllReduceFusionOp::RESIDUAL_RMS_NORM) { + FLASHINFER_CUDA_CALL(cudaLaunchKernelEx( + &kernelConfig, one_shot_all_reduce_norm_kernel, + params)); + } else { // fusionOp == AllReduceFusionOp::RESIDUAL_RMS_PREPOST_NORM + FLASHINFER_CUDA_CALL(cudaLaunchKernelEx( + &kernelConfig, one_shot_prenorm_all_reduce_norm_kernel, + params)); + } + } else { + cudaLaunchConfig_t kernelConfig = {0}; + kernelConfig.gridDim = cta_num; + kernelConfig.blockDim = cta_size; + kernelConfig.dynamicSmemBytes = smem_size; + kernelConfig.stream = stream; + + cudaLaunchAttribute attribute[1]; + attribute[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attribute[0].val.programmaticStreamSerializationAllowed = launch_with_pdl; + kernelConfig.attrs = attribute; + kernelConfig.numAttrs = 1; + + if (fusionOp == AllReduceFusionOp::RESIDUAL_RMS_NORM) { + FLASHINFER_CUDA_CALL(cudaLaunchKernelEx( + &kernelConfig, one_shot_all_reduce_norm_kernel, + params)); + } else { // fusionOp == AllReduceFusionOp::RESIDUAL_RMS_PREPOST_NORM + FLASHINFER_CUDA_CALL(cudaLaunchKernelEx( + &kernelConfig, one_shot_prenorm_all_reduce_norm_kernel, + params)); + } + } + } + return cudaSuccess; +} + +template +__global__ void lamport_initialize_kernel(T* buffer, size_t size) { + static constexpr uint32_t VEC_SIZE = 16 / sizeof(T); + for (size_t offset = (blockIdx.x * blockDim.x + threadIdx.x) * VEC_SIZE; offset < size; + offset += gridDim.x * blockDim.x * VEC_SIZE) { + set_neg_zero(&buffer[offset]); + } +} + +template +cudaError_t lamport_initialize_kernel_launcher(void* buffer, size_t size, cudaStream_t stream) { + static constexpr uint32_t VEC_SIZE = 16 / sizeof(T); + int block_size = 1024; + int grid_size = (size + 1024 * VEC_SIZE - 1) / (1024 * VEC_SIZE); + lamport_initialize_kernel + <<>>(reinterpret_cast(buffer), size); + auto status = cudaGetLastError(); + return status; +} +}; // namespace reduce_fusion + +template +static __global__ void oneShotAllReduceKernel(AllReduceParams params) { + // Suppose that two GPUs participate in the AR exchange, and we start four blocks. + // The message is partitioned into chunks as detailed below: + // message + // |-------------------| + // GPU 0 | B0 | B1 | B2 | B3 | + // GPU 1 | B0 | B1 | B2 | B3 | + // + // Here the step-by-step behavior of one block: + // 1. B0 copies the chunk it is responsible for, from local_input to shareable buffer + // 2. B0 on GPU 0 and B0 on GPU 1 wait for each other (block_barrier) + // 3. B0 on GPU 0 pull and sum the chunk from GPU 1, writes the result to local_output + // + // With COPY_INPUT == false, skip step 1. and use gpu_barrier instead of block barrier during + // step 2. We only to know if the other GPU as arrived at the AR kernel, that would mean that data + // is ready + // + // With PUSH_MODE, we consider that the shared buffer is of size: + // params.peer_comm_buffer_ptrs: [world_size, world_size, message_size] + // + // Here the step-by-step behavior of one block: + // 1. B0 push the chunk is it responsible for into all other GPUs: + // params.peer_comm_buffer_ptrs[:, local_gpu, B0 slice] + // 2. block sync so the block is shared by other GPUs + // 3. Reduce along second dimension params.peer_comm_buffer_ptrs[local_gpu, :, B0 slice] + + int const bidx = blockIdx.x; + int const tidx = threadIdx.x; + int const grid_size = gridDim.x; + + static constexpr uint32_t VEC_SIZE = 16 / sizeof(T); + + T const* local_input_buffer = reinterpret_cast(params.local_input_buffer_ptr); + T* local_shared_buffer = reinterpret_cast(params.peer_comm_buffer_ptrs[params.local_rank]); + T* local_output_buffer = reinterpret_cast(params.local_output_buffer_ptr); + + // Start and end offsets of the thread + size_t const chunk_start = bidx * params.elts_per_block + tidx * VEC_SIZE; + size_t const chunk_end = std::min((bidx + 1) * params.elts_per_block, params.elts_total); + + T* buffers[RANKS_PER_NODE]; +#pragma unroll + for (int ii = 0; ii < RANKS_PER_NODE; ++ii) { + // buffers[0] is always the local buffers. Helps load balancing reads. + int rank = (params.local_rank + ii) % RANKS_PER_NODE; + buffers[ii] = reinterpret_cast(params.peer_comm_buffer_ptrs[rank]); + } + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) && (__CUDA_ARCH__ < 1200)) + cudaGridDependencySynchronize(); +#endif + + if constexpr (PUSH_MODE || COPY_INPUT) { + // Copy from local buffer to shareable buffer + for (size_t iter_offset = chunk_start; iter_offset < chunk_end; + iter_offset += blockDim.x * VEC_SIZE) { + if constexpr (PUSH_MODE) { +#pragma unroll + for (int ii = 0; ii < RANKS_PER_NODE; ++ii) { + *reinterpret_cast( + &buffers[ii][params.local_rank * params.elts_total + iter_offset]) = + *reinterpret_cast(&local_input_buffer[iter_offset]); + } + } else { + *reinterpret_cast(&local_shared_buffer[iter_offset]) = + *reinterpret_cast(&local_input_buffer[iter_offset]); + } + } + + // wait for equivalent blocks of other GPUs to have copied data to their shareable buffer + block_barrier(params.peer_barrier_ptrs_in, params.barrier_flag, params.local_rank, + RANKS_PER_NODE, tidx, bidx, grid_size); + } else { + // In the non-copy case, we assume that once the kernel has been started, data is ready to be + // consumed + multi_gpu_barrier(params.peer_barrier_ptrs_in, params.barrier_flag, params.local_rank, + RANKS_PER_NODE, tidx, bidx); + } + + // Each block accumulates the values from the different GPUs on the same node. + for (size_t iter_offset = chunk_start; iter_offset < chunk_end; + iter_offset += blockDim.x * VEC_SIZE) { + // Iterate over the different ranks/devices on the node to load the values. + vec_t vals[RANKS_PER_NODE]; +#pragma unroll + for (int ii = 0; ii < RANKS_PER_NODE; ++ii) { + if constexpr (PUSH_MODE) { + vals[ii].load(&buffers[params.local_rank][ii * params.elts_total + iter_offset]); + } else { + vals[ii].load(&buffers[ii][iter_offset]); + } + } + + // Sum the values from the different ranks. + vec_t sums; + sums.fill(T(0)); +#pragma unroll + for (int rank = 0; rank < RANKS_PER_NODE; ++rank) { + // Always reduce from rank 0 to ensure stable reduce order. + int ii = (rank + RANKS_PER_NODE - params.local_rank) % RANKS_PER_NODE; + sums = vec_add(sums, vals[ii]); + } + // Store to the destination buffer. + sums.store(&local_output_buffer[iter_offset]); + } + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) && (__CUDA_ARCH__ < 1200)) + cudaTriggerProgrammaticLaunchCompletion(); +#endif +} + +template +static __global__ void __launch_bounds__(512, 1) twoShotAllReduceKernel(AllReduceParams params) { + // Suppose that two GPUs participate in the AR exchange, and we start two blocks. + // The message is partitioned into chunks as detailed below: + // message + // |-------------------| + // |--GPU 0--|--GPU 1--| (GPU responsibility parts) + // GPU 0 | B0 | B1 | B0 | B1 | + // GPU 1 | B0 | B1 | B0 | B1 | + // + // Here the step-by-step behavior of one block: + // 1. B0 copies all chunks is it responsible for, from local_input to shareable buffer + // 2. B0 on GPU 0 and B0 on GPU 1 wait for each other (block_barrier #0) + // 3. B0 on GPU 0 gather and sum the B0 chunks from GPU 1, that are in the GPU 0 responsibility + // part (the first half of the message, see GPU responsibility row above) + // 3bis. Likewise, B0 on GPU 1 copies and sum the chunks for GPU 0, + // where GPU 1 is responsible: the second half of the message. + // 4. B0 on GPU 0 and B0 on GPU 1 wait for each other (block_barrier #1) + // 5. B0 writes result to local_output. It gathers each chunk from its responsible GPU. + // For example, here it reads the first chunk from GPU 0 and second chunk from GPU 1. + // + // With COPY_INPUT == false, skip step 1. and use gpu_barrier instead of block barrier during + // step 2. We only to know if the other GPU as arrived at the AR kernel, that would mean that data + // is ready to be read. + // + // Note that compared to one-shot, one block (CTA) writes multiple input chunks and write multiple + // output chunks. However, it's only responsible for the summation of a single chunk. + // + // With PUSH_MODE, we consider that the shared buffer is of size: + // params.peer_comm_buffer_ptrs: [world_size, world_size, message_size / world_size] + // + // Here the step-by-step behavior of one block: + // 1. B0 push the chunks is it responsible for into the corresponding GPUs: + // params.peer_comm_buffer_ptrs[target_gpu, local_gpu, current B0 slice] + // 2. block sync so the blocks have been shared by other GPUs + // 3. Reduce along second dimension params.peer_comm_buffer_ptrs[local_gpu, :, B0 slice] + // 4. block barrier (corresponding blocks have finished reduction) + // 5. pull and write on local buffer, by reading params.peer_comm_buffer_ptrs[:, 0, B0 slice] + // (reduction result is + // written at index 0 of 2nd dim) + + int const bidx = blockIdx.x; + int const tidx = threadIdx.x; + int const grid_size = gridDim.x; + + static constexpr uint32_t VEC_SIZE = 16 / sizeof(T); + + T const* local_input_buffer = reinterpret_cast(params.local_input_buffer_ptr); + T* local_shared_buffer = reinterpret_cast(params.peer_comm_buffer_ptrs[params.local_rank]); + T* local_output_buffer = reinterpret_cast(params.local_output_buffer_ptr); + + size_t const chunk_start = bidx * params.elts_per_block + tidx * VEC_SIZE; + size_t const chunk_end = min(chunk_start + params.elts_per_block, params.elts_per_rank); + + T* buffers[RANKS_PER_NODE]; + int ranks[RANKS_PER_NODE]; +#pragma unroll + for (int ii = 0; ii < RANKS_PER_NODE; ++ii) { + // A mapping of the ranks to scatter reads as much as possible + int rank = (params.local_rank + ii) % RANKS_PER_NODE; + ranks[ii] = rank; + buffers[ii] = reinterpret_cast(params.peer_comm_buffer_ptrs[rank]); + } + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) && (__CUDA_ARCH__ < 1200)) + cudaGridDependencySynchronize(); +#endif + + if constexpr (PUSH_MODE || COPY_INPUT) { + // Copy all blocks from local buffer to shareable buffer + for (size_t local_offset = chunk_start; local_offset < chunk_end; + local_offset += blockDim.x * VEC_SIZE) { +#pragma unroll + for (int ii = 0; ii < RANKS_PER_NODE; ++ii) { + size_t offset_rank = ranks[ii] * params.elts_per_rank + local_offset; + if (offset_rank >= params.elts_total) { + continue; + } + + if constexpr (PUSH_MODE) { + *reinterpret_cast( + &buffers[ii][params.local_rank * params.elts_per_rank + local_offset]) = + *reinterpret_cast(&local_input_buffer[offset_rank]); + } else { + *reinterpret_cast(&local_shared_buffer[offset_rank]) = + *reinterpret_cast(&local_input_buffer[offset_rank]); + } + } + } + block_barrier(params.peer_barrier_ptrs_in, params.barrier_flag, params.local_rank, + RANKS_PER_NODE, tidx, bidx, grid_size); + } else { + // In the non-copy case, we assume that once the kernel has been started, data is ready to be + // consumed + multi_gpu_barrier(params.peer_barrier_ptrs_in, params.barrier_flag, params.local_rank, + RANKS_PER_NODE, tidx, bidx); + } + + // Each block accumulates the values from the different GPUs on the same node. + for (size_t local_offset = chunk_start; local_offset < chunk_end; + local_offset += blockDim.x * VEC_SIZE) { + size_t const responsible_block_offset = local_offset + params.rank_offset; + + // Iterate over the different ranks/devices on the node to load the values. + vec_t vals[RANKS_PER_NODE]; +#pragma unroll + for (int ii = 0; ii < RANKS_PER_NODE; ++ii) { + if constexpr (PUSH_MODE) { + vals[ii].load(&local_shared_buffer[ii * params.elts_per_rank + local_offset]); + } else { + vals[ii].load(&buffers[ii][responsible_block_offset]); + } + } + + // Sum the values from the different ranks. + vec_t sums; + sums.fill(T(0)); +#pragma unroll + for (int rank = 0; rank < RANKS_PER_NODE; ++rank) { + // Always reduce from rank 0 to ensure stable reduce order. + int ii = (rank + RANKS_PER_NODE - params.local_rank) % RANKS_PER_NODE; + sums = vec_add(sums, vals[ii]); + } + + // Store to the local buffer. + if constexpr (PUSH_MODE) { + sums.store(&local_shared_buffer[local_offset]); + } else { + sums.store(&local_shared_buffer[responsible_block_offset]); + } + } + + block_barrier(params.peer_barrier_ptrs_out, params.barrier_flag, params.local_rank, + RANKS_PER_NODE, tidx, bidx, grid_size); + + // Gather all needed elts from other intra-node ranks + for (size_t local_offset = chunk_start; local_offset < chunk_end; + local_offset += blockDim.x * VEC_SIZE) { +#pragma unroll + for (int ii = 0; ii < RANKS_PER_NODE; ++ii) { + // use round-robin gathering from other ranks + size_t offset_rank = ranks[ii] * params.elts_per_rank + local_offset; + if (offset_rank >= params.elts_total) { + continue; + } + vec_t sums, residual_vec, bias_vec; + if constexpr (Bias) { + bias_vec.load(reinterpret_cast(params.fusion_params.bias_buffer) + + offset_rank % params.fusion_params.hidden_size); + } + if constexpr (Residual) { + residual_vec.load(reinterpret_cast(params.fusion_params.residual_buffer) + + offset_rank); + } + if constexpr (PUSH_MODE) { + *reinterpret_cast(&local_output_buffer[offset_rank]) = + *reinterpret_cast(&buffers[ii][local_offset]); + sums.load(&buffers[ii][local_offset]); + } else { + *reinterpret_cast(&local_output_buffer[offset_rank]) = + *reinterpret_cast(&buffers[ii][offset_rank]); + sums.load(&buffers[ii][offset_rank]); + } + if constexpr (Bias) { + sums = vec_add(sums, bias_vec); + } + if constexpr (Residual) { + sums = vec_add(sums, residual_vec); + } + sums.store(&local_output_buffer[offset_rank]); + } + } + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) && (__CUDA_ARCH__ < 1200)) + cudaTriggerProgrammaticLaunchCompletion(); +#endif +} + +template +bool configurationSupported(AllReduceStrategyType algo, size_t msg_size, size_t n_ranks) { + size_t elts_per_thread = 16 / sizeof(T); + int const msg_align = + (algo == AllReduceStrategyType::TWOSHOT) ? n_ranks * elts_per_thread : elts_per_thread; + bool supported_algo = + (algo == AllReduceStrategyType::ONESHOT || algo == AllReduceStrategyType::TWOSHOT); + return supported_algo && (msg_size % msg_align == 0); +} + +template +std::tuple kernelLaunchConfig(AllReduceStrategyType algo, AllReduceParams& params, + size_t elts_per_thread) { + int blocks_per_grid = 1, threads_per_block = DEFAULT_BLOCK_SIZE; + + switch (algo) { + case AllReduceStrategyType::ONESHOT: { + FLASHINFER_CHECK(params.elts_total % elts_per_thread == 0, + "hidden_size must be a multiple of ", elts_per_thread); + size_t const total_threads = round_up(params.elts_total / elts_per_thread, WARP_SIZE); + threads_per_block = std::min(DEFAULT_BLOCK_SIZE, total_threads); + blocks_per_grid = std::min(static_cast(MAX_ALL_REDUCE_BLOCKS), + ceil_div(total_threads, threads_per_block)); + params.elts_per_block = + round_up(ceil_div(params.elts_total, blocks_per_grid), elts_per_thread); + break; + } + case AllReduceStrategyType::TWOSHOT: { + FLASHINFER_CHECK(params.elts_total % (elts_per_thread * params.ranks_per_node) == 0, + "hidden_size must be a multiple of ", + elts_per_thread * params.ranks_per_node); + size_t const total_threads = + round_up(params.elts_total / (elts_per_thread * params.ranks_per_node), WARP_SIZE); + + /* + threads_per_block = std::min(DEFAULT_BLOCK_SIZE, total_threads); + blocks_per_grid = std::min(static_cast(MAX_ALL_REDUCE_BLOCKS), ceil_div(total_threads, + threads_per_block)); + */ + while (total_threads % blocks_per_grid != 0 || + total_threads / blocks_per_grid > DEFAULT_BLOCK_SIZE) { + blocks_per_grid += 1; + } + + threads_per_block = total_threads / blocks_per_grid; + + // NOTE: need to adjust here + if (blocks_per_grid > MAX_ALL_REDUCE_BLOCKS) { + size_t iter_factor = 1; + while (blocks_per_grid / iter_factor > MAX_ALL_REDUCE_BLOCKS || + blocks_per_grid % iter_factor) { + iter_factor += 1; + } + blocks_per_grid /= iter_factor; + } + params.elts_per_rank = params.elts_total / params.ranks_per_node; + params.rank_offset = params.local_rank * params.elts_per_rank; + params.elts_per_block = + round_up(ceil_div(params.elts_per_rank, blocks_per_grid), elts_per_thread); + break; + } + default: + FLASHINFER_ERROR("Algorithm not supported here."); + } + + return std::make_tuple(blocks_per_grid, threads_per_block); +} + +template +cudaError_t AllReduceNormKernelLaunch(AllReduceStrategyType algo, AllReduceFusionOp fusionOp, + AllReduceParams& params, bool launch_with_pdl, + cudaStream_t stream) { + FLASHINFER_CHECK((fusionOp == AllReduceFusionOp::RESIDUAL_RMS_NORM || + fusionOp == AllReduceFusionOp::RESIDUAL_RMS_PREPOST_NORM), + "Unsupported AllReduceFusionOp: %d", static_cast(fusionOp)); + if (algo == AllReduceStrategyType::ONESHOT) { + return reduce_fusion::one_shot_all_reduce_norm_kernel_launcher( + params, fusionOp, launch_with_pdl, stream); + } else { + FLASHINFER_CHECK(!(USE_MEMCPY && PUSH_MODE), "Memcpy cannot be used with PUSH_MODE."); + size_t elts_per_thread = 16 / sizeof(T); + auto [blocks_per_grid, threads_per_block] = + kernelLaunchConfig(algo, params, elts_per_thread); + if (USE_MEMCPY) { + cudaMemcpyAsync(params.peer_comm_buffer_ptrs[params.local_rank], + params.local_input_buffer_ptr, params.elts_total * sizeof(T), + cudaMemcpyDeviceToDevice, stream); + } + auto output_ptr = params.local_output_buffer_ptr; + params.local_output_buffer_ptr = params.fusion_params.intermediate_buffer; + + cudaLaunchConfig_t kernelConfig = {0}; + kernelConfig.gridDim = blocks_per_grid; + kernelConfig.blockDim = threads_per_block; + kernelConfig.dynamicSmemBytes = 0; + kernelConfig.stream = stream; + + cudaLaunchAttribute attribute[1]; + attribute[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attribute[0].val.programmaticStreamSerializationAllowed = launch_with_pdl; + kernelConfig.attrs = attribute; + kernelConfig.numAttrs = 1; + + FLASHINFER_CUDA_CALL(cudaLaunchKernelEx( + &kernelConfig, + twoShotAllReduceKernel, params)); + + params.local_output_buffer_ptr = output_ptr; + return reduce_fusion::rms_norm_kernel_launcher( + params, fusionOp, launch_with_pdl, stream); + } + return cudaSuccess; +} + +template +cudaError_t AllReduceNormDispatch(AllReduceStrategyType algo, AllReduceFusionOp fusionOp, + AllReduceParams& params, bool launch_with_pdl, + cudaStream_t stream) { + if (params.fusion_params.bias_buffer && params.fusion_params.weight_buffer) { + return AllReduceNormKernelLaunch( + algo, fusionOp, params, launch_with_pdl, stream); + } else if (params.fusion_params.bias_buffer && !params.fusion_params.weight_buffer) { + return AllReduceNormKernelLaunch( + algo, fusionOp, params, launch_with_pdl, stream); + } else if (!params.fusion_params.bias_buffer && params.fusion_params.weight_buffer) { + return AllReduceNormKernelLaunch( + algo, fusionOp, params, launch_with_pdl, stream); + } else { + return AllReduceNormKernelLaunch( + algo, fusionOp, params, launch_with_pdl, stream); + } + return cudaSuccess; +} + +template +cudaError_t AllReduceDispatch(AllReduceStrategyType algo, AllReduceFusionOp fusionOp, + AllReduceParams& params, bool launch_with_pdl, + cudaStream_t stream) { + FLASHINFER_CHECK(fusionOp == AllReduceFusionOp::NONE, "Unsupported AllReduceFusionOp: %d", + static_cast(fusionOp)); + FLASHINFER_CHECK(!(USE_MEMCPY && PUSH_MODE), "Memcpy cannot be used with PUSH_MODE."); + size_t elts_per_thread = 16 / sizeof(T); + auto [blocks_per_grid, threads_per_block] = kernelLaunchConfig(algo, params, elts_per_thread); + if (USE_MEMCPY) { + cudaMemcpyAsync(params.peer_comm_buffer_ptrs[params.local_rank], params.local_input_buffer_ptr, + params.elts_total * sizeof(T), cudaMemcpyDeviceToDevice, stream); + } + if (algo == AllReduceStrategyType::ONESHOT) { + auto* kernel_instance = &oneShotAllReduceKernel; + cudaLaunchConfig_t config; + config.gridDim = blocks_per_grid; + config.blockDim = threads_per_block; + config.dynamicSmemBytes = 0; + config.stream = stream; + cudaLaunchAttribute attribute[1]; + attribute[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attribute[0].val.programmaticStreamSerializationAllowed = launch_with_pdl; + config.attrs = attribute; + config.numAttrs = 1; + cudaLaunchKernelEx(&config, kernel_instance, params); + } else { + auto* kernel_instance = &twoShotAllReduceKernel; + cudaLaunchConfig_t config; + config.gridDim = blocks_per_grid; + config.blockDim = threads_per_block; + config.dynamicSmemBytes = 0; + config.stream = stream; + cudaLaunchAttribute attribute[1]; + attribute[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attribute[0].val.programmaticStreamSerializationAllowed = launch_with_pdl; + config.attrs = attribute; + config.numAttrs = 1; + cudaLaunchKernelEx(&config, kernel_instance, params); + } + return cudaSuccess; +} + +template +cudaError_t AllReduceDispatchMemcpy(AllReduceStrategyType algo, AllReduceFusionOp fusionOp, + AllReduceParams& params, bool launch_with_pdl, + cudaStream_t stream) { + if (fusionOp == AllReduceFusionOp::NONE) { + FLASHINFER_LOG_DEBUG("AllReduceDispatch enabled"); + return AllReduceDispatch(algo, fusionOp, params, + launch_with_pdl, stream); + } else { + FLASHINFER_LOG_DEBUG("AllReduceNormDispatch enabled"); + return AllReduceNormDispatch(algo, fusionOp, params, + launch_with_pdl, stream); + } +} + +template +cudaError_t AllReduceDispatchPushMode(AllReduceStrategyType algo, AllReduceStrategyConfig config, + AllReduceFusionOp fusionOp, AllReduceParams& params, + bool launch_with_pdl, cudaStream_t stream) { + if (static_cast>(config) & + static_cast>( + AllReduceStrategyConfig::USE_MEMCPY)) { + return AllReduceDispatchMemcpy(algo, fusionOp, params, + launch_with_pdl, stream); + } else { + return AllReduceDispatchMemcpy(algo, fusionOp, params, + launch_with_pdl, stream); + } +} + +template //, bool USE_MEMCPY = false, bool PUSH_MODE = false> +cudaError_t AllReduceDispatchRanksPerNode(AllReduceStrategyType algo, + AllReduceStrategyConfig config, + AllReduceFusionOp fusionOp, AllReduceParams& params, + bool launch_with_pdl, cudaStream_t stream) { + if (static_cast>(config) & + static_cast>( + AllReduceStrategyConfig::PUSH_MODE)) { + return AllReduceDispatchPushMode(algo, config, fusionOp, params, + launch_with_pdl, stream); + } else { + return AllReduceDispatchPushMode(algo, config, fusionOp, params, + launch_with_pdl, stream); + } +} + +template +cudaError_t AllReduceDispatchType(AllReduceParams& params, AllReduceStrategyType strat, + AllReduceStrategyConfig config, AllReduceFusionOp fusionOp, + bool launch_with_pdl, cudaStream_t stream) { + switch (params.ranks_per_node) { + case 2: + return AllReduceDispatchRanksPerNode(strat, config, fusionOp, params, + launch_with_pdl, stream); + case 4: + return AllReduceDispatchRanksPerNode(strat, config, fusionOp, params, launch_with_pdl, + stream); + case 6: + return AllReduceDispatchRanksPerNode(strat, config, fusionOp, params, launch_with_pdl, + stream); + case 8: + return AllReduceDispatchRanksPerNode(strat, config, fusionOp, params, launch_with_pdl, + stream); + case 16: + return AllReduceDispatchRanksPerNode(strat, config, fusionOp, params, launch_with_pdl, + stream); + default: + FLASHINFER_ERROR("Custom all reduce only supported on {2, 4, 6, 8, 16} GPUs per node."); + } + return cudaSuccess; +} + +template +cudaError_t customAllReduce(AllReduceParams& params, AllReduceStrategyType strat, + AllReduceStrategyConfig config, AllReduceFusionOp fusionOp, + bool launch_with_pdl, cudaStream_t stream) { + FLASHINFER_CHECK(configurationSupported(strat, params.elts_total, params.ranks_per_node), + "Custom all-reduce configuration unsupported"); + + return AllReduceDispatchType(params, strat, config, fusionOp, launch_with_pdl, stream); +} + +template +cudaError_t lamportInitialize(void* buffer, size_t size, cudaStream_t stream) { + if (size == 0) { + return cudaSuccess; + } + FLASHINFER_LOG_INFO("lamportInitialize start: buffer: {}, size: {}", buffer, size); + return reduce_fusion::lamport_initialize_kernel_launcher(buffer, size, stream); +} + +// lamport: 3 buffers for synchronization +template +cudaError_t lamportInitializeAll(void* buffer_0, void* buffer_1, void* buffer_2, size_t size, + cudaStream_t stream) { + auto status = lamportInitialize(buffer_0, size / sizeof(T), stream); + FLASHINFER_CHECK(status == cudaSuccess, "lamportInitialize failed with error code " + + std::string(cudaGetErrorString(status))); + + status = lamportInitialize(buffer_1, size / sizeof(T), stream); + FLASHINFER_CHECK(status == cudaSuccess, "lamportInitialize failed with error code " + + std::string(cudaGetErrorString(status))); + + status = lamportInitialize(buffer_2, size / sizeof(T), stream); + FLASHINFER_CHECK(status == cudaSuccess, "lamportInitialize failed with error code " + + std::string(cudaGetErrorString(status))); + cudaDeviceSynchronize(); + return cudaSuccess; +} + +} // namespace trtllm_allreduce +} // namespace flashinfer diff --git a/include/flashinfer/comm/trtllm_allreduce_fusion.cuh b/include/flashinfer/comm/trtllm_allreduce_fusion.cuh new file mode 100644 index 0000000000000000000000000000000000000000..dc48372cc757e973102e3ffe188f73c1abb18cae --- /dev/null +++ b/include/flashinfer/comm/trtllm_allreduce_fusion.cuh @@ -0,0 +1,1471 @@ +#include +#include +#include +#include + +#include +#include +#include + +#include "../exception.h" +#include "../logging.h" +#include "../utils.cuh" +#include "../vec_dtypes.cuh" + +namespace flashinfer { + +namespace trtllm_allreduce_fusion { + +enum class FP4QuantizationSFLayout { + // Block scale factors are stored in swizzled layout for cutlass FP4 kernel. Scale factor + // blocks are organized in 512-byte blocks in global memory, with each block having 128x4 FP8 + // values. The SF matrix dimensions are therefore padded - rows to the nearest multiple of 128 and + // columns to the nearest multiple of 4. + // + // The scale factor block rows map to data block rows in an interleaved pattern: + // For a scale factor row 'i', it maps to data block row: (i % 4) * 32 + (i / 4) + // Column 'j' in the scale factor block corresponds to scaling the j-th block in the data tensor. + // + // Please refer to https://nvbugs/4165523 for more details about the swizzled layout. + SWIZZLED, + // Block scale factors are stored in linear layout (row-major). This is used in some trtllm-gen + // kernels standard. + LINEAR +}; + +namespace details { + +static constexpr int CVT_FP4_ELTS_PER_THREAD = 8; +static constexpr int CVT_FP4_SF_VEC_SIZE = 16; +static constexpr int kBytesPerAccess = 16; +static constexpr int kOneShotMaxToken = 128; +static constexpr int kBarrierFlagCount = 256; + +} // namespace details + +namespace maths { +// // ============================== Cast ============================== +template +__device__ inline T_OUT cuda_cast(T_IN val) { + return val; +} + +template <> +__device__ inline float2 cuda_cast(int2 val) { + return make_float2(val.x, val.y); +} + +template <> +__device__ inline float2 cuda_cast(float val) { + return make_float2(val, val); +} + +template <> +__device__ inline float2 cuda_cast(half2 val) { + return __half22float2(val); +} + +template <> +__device__ inline half2 cuda_cast(float2 val) { + return __float22half2_rn(val); +} + +template <> +__device__ inline half2 cuda_cast(float val) { + return __float2half2_rn(val); +} + +template <> +__device__ inline half2 cuda_cast(half val) { + return __half2half2(val); +} + +template <> +__device__ inline int8_t cuda_cast(half val) { + union { + int8_t int8[2]; + int16_t int16; + }; + + union { + half fp16; + int16_t int16_in; + }; + + fp16 = val; + asm volatile("cvt.rni.sat.s8.f16 %0, %1;" : "=h"(int16) : "h"(int16_in)); + return int8[0]; +} + +template <> +__device__ inline int16_t cuda_cast(half2 val) { + union { + int8_t int8[2]; + int16_t int16; + }; + + int8[0] = cuda_cast(val.x); + int8[1] = cuda_cast(val.y); + return int16; +} + +template <> +__device__ inline int8_t cuda_cast(float val) { + union { + int8_t int8[2]; + int16_t int16; + }; + + asm volatile("cvt.rni.sat.s8.f32 %0, %1;" : "=h"(int16) : "f"(val)); + return int8[0]; +} + +template <> +__device__ inline int16_t cuda_cast(float2 val) { + union { + int8_t int8[2]; + int16_t int16; + }; + + int8[0] = cuda_cast(val.x); + int8[1] = cuda_cast(val.y); + return int16; +} + +template <> +__device__ inline half2 cuda_cast(int16_t val) { + union { + int8_t int8[2]; + int16_t int16; + }; + + int16 = val; + return make_half2(int8[0], int8[1]); +} + +template <> +__device__ inline float2 cuda_cast(int16_t val) { + union { + int8_t int8[2]; + int16_t int16; + }; + + int16 = val; + return make_float2(int8[0], int8[1]); +} + +template <> +__device__ inline __nv_bfloat16 cuda_cast(int32_t val) { + return static_cast(val); +} + +template <> +__device__ inline __nv_bfloat16 cuda_cast(int8_t val) { + return static_cast(val); +} + +template <> +__device__ inline int8_t cuda_cast(__nv_bfloat16 val) { + return static_cast(val); +} + +template <> +__device__ inline float cuda_cast(__nv_bfloat16 val) { + return __bfloat162float(val); +} + +inline __device__ float2 bf1622float2(const __nv_bfloat162 val) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + float2 f_val; + f_val.x = __low2float(val); + f_val.y = __high2float(val); + return f_val; +#else + return __bfloat1622float2(val); +#endif +} + +template <> +__device__ inline float2 cuda_cast(__nv_bfloat162 val) { + return bf1622float2(val); +} + +template <> +__device__ inline half cuda_cast(__nv_bfloat16 val) { + return __float2half(__bfloat162float(val)); +} + +inline __device__ int16_t bf1622int16(__nv_bfloat162 val) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + float2 f_val; + f_val.x = max(min(__low2float(val), 127.f), -128.f); + f_val.y = max(min(__high2float(val), 127.f), -128.f); + + union { + int8_t int8[2]; + int16_t int16; + }; + + int8[0] = static_cast(static_cast(f_val.x)); + int8[1] = static_cast(static_cast(f_val.y)); + return int16; +#else + val = __hmin2(val, make_bfloat162(127., 127.)); + val = __hmax2(val, make_bfloat162(-128., -128.)); + + union { + int8_t int8[2]; + int16_t int16; + }; + + int8[0] = static_cast(static_cast(val.x)); + int8[1] = static_cast(static_cast(val.y)); + return int16; +#endif +} + +template <> +__device__ inline int16_t cuda_cast(__nv_bfloat162 val) { + return bf1622int16(val); +} + +template <> +__device__ inline __nv_bfloat16 cuda_cast<__nv_bfloat16, float>(float val) { + return __float2bfloat16(val); +} + +template <> +__device__ inline __nv_bfloat16 cuda_cast<__nv_bfloat16, half>(half val) { + return __float2bfloat16(__half2float(val)); +} + +inline __device__ __nv_bfloat162 bf162bf162(const __nv_bfloat16 val) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + __nv_bfloat162 val2; + val2.x = val; + val2.y = val; + return val2; +#else + return __bfloat162bfloat162(val); +#endif +} + +template <> +__device__ inline __nv_bfloat162 cuda_cast<__nv_bfloat162, __nv_bfloat16>(__nv_bfloat16 val) { + return bf162bf162(val); +} + +template <> +__device__ inline __nv_bfloat162 cuda_cast<__nv_bfloat162, float>(float val) { + return __float2bfloat162_rn(val); +} + +inline __device__ __nv_bfloat162 float22bf162(const float2 val) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + return __floats2bfloat162_rn(val.x, val.y); +#else + return __float22bfloat162_rn(val); +#endif +} + +template <> +__device__ inline __nv_bfloat162 cuda_cast<__nv_bfloat162, float2>(float2 val) { + return float22bf162(val); +} + +template <> +__device__ inline __nv_bfloat162 cuda_cast<__nv_bfloat162, int16_t>(int16_t val) { + union { + int8_t int8[2]; + int16_t int16; + }; + + int16 = val; + __nv_bfloat162 res; + res.x = cuda_cast<__nv_bfloat16>(int8[0]); + res.y = cuda_cast<__nv_bfloat16>(int8[1]); + return res; +} + +template <> +__device__ inline __nv_bfloat162 cuda_cast<__nv_bfloat162, half2>(half2 val) { + return float22bf162(__half22float2(val)); +} + +// // ============================== Abs ============================== +template +__device__ inline T cuda_abs(T val) { + assert(false); + return {}; +} + +template <> +__device__ inline float cuda_abs(float val) { + return fabs(val); +} + +template <> +__device__ inline float2 cuda_abs(float2 val) { + return make_float2(fabs(val.x), fabs(val.y)); +} + +template <> +__device__ inline half cuda_abs(half val) { + return __habs(val); +} + +template <> +__device__ inline half2 cuda_abs(half2 val) { + return __habs2(val); +} + +#if __CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__) +template <> +__device__ inline __nv_bfloat16 cuda_abs(__nv_bfloat16 val) { + return __habs(val); +} + +template <> +__device__ inline __nv_bfloat162 cuda_abs(__nv_bfloat162 val) { + return __habs2(val); +} +#endif + +// // ============================== Max ============================== +template +__device__ inline To cuda_max(Ti val) { + return cuda_cast(val); +}; + +template <> +__device__ inline float cuda_max(float2 val) { + return fmaxf(val.x, val.y); +} + +template <> +__device__ inline half cuda_max(half2 val) { + return __hmax(val.x, val.y); +} + +template <> +__device__ inline __nv_bfloat16 cuda_max(__nv_bfloat162 val) { +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)) + return __hmax(val.x, val.y); +#else + assert(0); + asm volatile("brkpt;\n" ::); + return __nv_bfloat16(0); +#endif +} + +// Binary maximum: compute the max of two values. +template +__device__ inline T cuda_max(T val1, T val2) { + return (val1 > val2) ? val1 : val2; +} + +template <> +__device__ inline float2 cuda_max(float2 val1, float2 val2) { + float2 out; + out.x = fmaxf(val1.x, val2.x); + out.y = fmaxf(val1.y, val2.y); + return out; +} + +template <> +__device__ inline half2 cuda_max(half2 val1, half2 val2) { + return __hmax2(val1, val2); +} + +template <> +__device__ inline __nv_bfloat162 cuda_max(__nv_bfloat162 val1, __nv_bfloat162 val2) { + return __hmax2(val1, val2); +} + +// // ============================== Reciprocal ============================== +// Fast reciprocal. +inline __device__ float reciprocal_approximate_ftz(float a) { + float b; + asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a)); + return b; +} +} // namespace maths + +namespace utils { + +#define FINAL_MASK 0xffffffff + +template +__inline__ __device__ T warpReduceSumV2(T* val) { +#pragma unroll + for (int i = 0; i < NUM; i++) { +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) + val[i] += __shfl_xor_sync(FINAL_MASK, val[i], mask, 32); + } + return (T)(0.0f); +} + +template +__inline__ __device__ T blockReduceSumV2(T* val) { + static __shared__ T shared[NUM][33]; + int lane = threadIdx.x & 0x1f; + int wid = threadIdx.x >> 5; + + warpReduceSumV2(val); + + if (lane == 0) { +#pragma unroll + for (int i = 0; i < NUM; i++) { + shared[i][wid] = val[i]; + } + } + + __syncthreads(); + + bool is_mask = threadIdx.x < (blockDim.x / 32.f); +#pragma unroll + for (int i = 0; i < NUM; i++) { + val[i] = is_mask ? shared[i][lane] : (T)(0.0f); + } + warpReduceSumV2(val); + return (T)0.0f; +} + +inline int getSMVersion() { + int device{-1}; + FLASHINFER_CUDA_CALL(cudaGetDevice(&device)); + int sm_major = 0; + int sm_minor = 0; + FLASHINFER_CUDA_CALL( + cudaDeviceGetAttribute(&sm_major, cudaDevAttrComputeCapabilityMajor, device)); + FLASHINFER_CUDA_CALL( + cudaDeviceGetAttribute(&sm_minor, cudaDevAttrComputeCapabilityMinor, device)); + return sm_major * 10 + sm_minor; +} + +inline __device__ int64_t get_sf_out_offset_128x4(std::optional batchIdx, int mIdx, int kIdx, + std::optional numRows, int numCols) { + // SF layout [numMTiles, numKTiles, 32 (mTile), 4 (mTile), 4(kTile)] + // --> index [mTileIdx, kTileIdx, outerMIdx, innerMIdx, innerKIdx] + + // batched tensor + // SF layout [numBTiles, numMTiles, numKTiles, 32 (mTile), 4 (mTile), 4(kTile)] + // --> index [bTileIdx, mTileIdx, kTileIdx, outerMIdx, innerMIdx, innerKIdx] + + int32_t innerKIdx = (kIdx % 4); + int64_t innerKStride = 1; + + int32_t innerMIdx = (mIdx % (32 * 4)) / 32; + int64_t innerMStride = 4 * innerKStride; // 4 + + // M tile layout [32, 4] is column-major. + int32_t outerMIdx = (mIdx % 32); + int64_t outerMStride = 4 * innerMStride; // 16 + + int32_t kTileIdx = (kIdx / 4); + int64_t kTileStride = 32 * outerMStride; // 512 + + // SF vector size 16. We round the "numCols" up to a multiple of 64. + int factor = details::CVT_FP4_SF_VEC_SIZE * 4; + int32_t numKTiles = (numCols + factor - 1) / factor; + int32_t mTileIdx = mIdx / (32 * 4); + int64_t mTileStride = numKTiles * kTileStride; + + // Each SF block has 128 rows so pad rows to the multiple of 128. + int32_t numMTiles = (numRows.value_or(0) + 128 - 1) / 128; + int64_t bTileStride = numMTiles * mTileStride; + + // Compute the global offset. + int64_t SFOffset = batchIdx.value_or(0) * bTileStride + mTileIdx * mTileStride + + kTileIdx * kTileStride + outerMIdx * outerMStride + innerMIdx * innerMStride + + innerKIdx * innerKStride; + + return SFOffset; +} + +template +__device__ uint8_t* cvt_quant_to_fp4_get_sf_out_offset(std::optional batchIdx, int rowIdx, + int colIdx, std::optional numRows, + int numCols, SFType* SFout, + FP4QuantizationSFLayout layout) { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + static_assert(CVT_FP4_NUM_THREADS_PER_SF == 1 || CVT_FP4_NUM_THREADS_PER_SF == 2); + + // One pair of threads write one SF to global memory. + // TODO: stage through smem for packed STG.32 + // is it better than STG.8 from 4 threads ? + if (threadIdx.x % CVT_FP4_NUM_THREADS_PER_SF == 0) { + if (layout == FP4QuantizationSFLayout::SWIZZLED) { + // SF vector index (16 elements share one SF in the K dimension). + // numRows and numCols are unpadded. + int32_t kIdx = colIdx / CVT_FP4_NUM_THREADS_PER_SF; + int32_t mIdx = rowIdx; + + auto SFOffset = get_sf_out_offset_128x4(batchIdx, mIdx, kIdx, numRows, numCols); + return reinterpret_cast(SFout) + SFOffset; + } else if (layout == FP4QuantizationSFLayout::LINEAR) { + // Linear row-major layout, no padding required. + int32_t KTileIdx = colIdx / CVT_FP4_NUM_THREADS_PER_SF; + + int32_t numKTiles = numCols / details::CVT_FP4_SF_VEC_SIZE; + int64_t mTileStride = numKTiles; + + int64_t BTileStride = numRows.value_or(0) * mTileStride; + + int64_t SFOffset = batchIdx.value_or(0) * BTileStride + rowIdx * mTileStride + KTileIdx; + return reinterpret_cast(SFout) + SFOffset; + } else { + return nullptr; + } + } +#endif + return nullptr; +} + +__forceinline__ __device__ uint32_t pack_bytes(uint8_t c0, uint8_t c1, uint8_t c2, uint8_t c3) { + uint32_t val0 = c0; + uint32_t val1 = c1; + uint32_t val2 = c2; + uint32_t val3 = c3; + + return (val3 << 24) | (val2 << 16) | (val1 << 8) | val0; +} + +// Convert 8 float32 values into 8 e2m1 values (represented as one uint32_t). +// NOTE: bypass sm_100 requirement by __nv_cvt_float2_to_fp4x2 +inline __device__ uint32_t fp32_vec_to_e2m1(float (&array)[8]) { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + uint32_t val; + asm volatile( + "{\n" + ".reg .b8 byte0;\n" + ".reg .b8 byte1;\n" + ".reg .b8 byte2;\n" + ".reg .b8 byte3;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte0, %2, %1;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte1, %4, %3;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte2, %6, %5;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte3, %8, %7;\n" + "mov.b32 %0, {byte0, byte1, byte2, byte3};\n" + "}" + : "=r"(val) + : "f"(array[0]), "f"(array[1]), "f"(array[2]), "f"(array[3]), "f"(array[4]), "f"(array[5]), + "f"(array[6]), "f"(array[7])); + return val; +#else + uint32_t val; + __nv_fp4x2_storage_t vals[4]; +#pragma unroll + for (int i = 0; i < 4; i++) { + vals[i] = __nv_cvt_float2_to_fp4x2(*(((float2*)array) + i), __NV_E2M1, cudaRoundNearest); + } + val = pack_bytes(vals[0], vals[1], vals[2], vals[3]); + return val; +#endif +} + +// Convert 4 float2 values into 8 e2m1 values (represented as one uint32_t). +inline __device__ uint32_t fp32_vec_to_e2m1(float2 (&array)[4]) { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + uint32_t val; + asm volatile( + "{\n" + ".reg .b8 byte0;\n" + ".reg .b8 byte1;\n" + ".reg .b8 byte2;\n" + ".reg .b8 byte3;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte0, %2, %1;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte1, %4, %3;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte2, %6, %5;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte3, %8, %7;\n" + "mov.b32 %0, {byte0, byte1, byte2, byte3};\n" + "}" + : "=r"(val) + : "f"(array[0].x), "f"(array[0].y), "f"(array[1].x), "f"(array[1].y), "f"(array[2].x), + "f"(array[2].y), "f"(array[3].x), "f"(array[3].y)); + return val; +#else + uint32_t val; + __nv_fp4x2_storage_t vals[4]; +#pragma unroll + for (int i = 0; i < 4; i++) { + vals[i] = __nv_cvt_float2_to_fp4x2(array[i], __NV_E2M1, cudaRoundNearest); + } + val = pack_bytes(vals[0], vals[1], vals[2], vals[3]); + return val; +#endif +} + +// Quantizes the provided PackedVec into the uint32_t output +template +__device__ uint32_t cvt_warp_fp16_to_fp4(vec_t& vec, float SFScaleVal, + uint8_t* SFout) { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + // Get absolute maximum values among the local 8 values. + auto localMax = maths::cuda_abs(get_vec2_element(vec, 0)); + +#pragma unroll + for (int i = 1; i < details::CVT_FP4_ELTS_PER_THREAD / 2; i++) { + localMax = maths::cuda_max(localMax, maths::cuda_abs(get_vec2_element(vec, i))); + } + + // Get the absolute maximum among all 16 values (two threads). + localMax = maths::cuda_max(__shfl_xor_sync(uint32_t(-1), localMax, 1), localMax); + // Get the final absolute maximum values. + float vecMax = float(maths::cuda_max(localMax.x, localMax.y)); + + // Get the SF (max value of the vector / max value of e2m1). + // maximum value of e2m1 = 6.0. + // TODO: use half as compute data type. + float SFValue = SFScaleVal * (vecMax * maths::reciprocal_approximate_ftz(6.0f)); + // 8 bits representation of the SF. + uint8_t fp8SFVal; + // Write the SF to global memory (STG.8). + if constexpr (UE8M0_SF) { +#if (__CUDACC_VER_MAJOR__ * 10000 + __CUDACC_VER_MINOR__ * 100 >= 120800) + __nv_fp8_e8m0 tmp; + tmp.__x = __nv_cvt_float_to_e8m0(SFValue, __NV_SATFINITE, cudaRoundPosInf); + SFValue = static_cast(tmp); + fp8SFVal = tmp.__x; +#else +#error "FP8 E8M0 support requires CUDA 12.8 or newer." +#endif + } else { + // Here SFValue is always positive, so E4M3 is the same as UE4M3. + __nv_fp8_e4m3 tmp = __nv_fp8_e4m3(SFValue); + fp8SFVal = tmp.__x; + SFValue = static_cast(tmp); + } + // Get the output scale. + // Recipe: final_scale = reciprocal(fp32(fp8(SFValue * SFScaleVal))) * reciprocal(SFScaleVal)) + float outputScale = SFValue != 0 ? maths::reciprocal_approximate_ftz( + SFValue * maths::reciprocal_approximate_ftz(SFScaleVal)) + : 0.0f; + + if (SFout) { + // Write the SF to global memory (STG.8). + *SFout = fp8SFVal; + } + + // Convert the input to float. + float2 fp2Vals[details::CVT_FP4_ELTS_PER_THREAD / 2]; + +#pragma unroll + for (int i = 0; i < details::CVT_FP4_ELTS_PER_THREAD / 2; i++) { + if constexpr (std::is_same_v) { + fp2Vals[i] = __half22float2(get_vec2_element(vec, i)); + } else { + fp2Vals[i] = __bfloat1622float2(get_vec2_element(vec, i)); + } + fp2Vals[i].x *= outputScale; + fp2Vals[i].y *= outputScale; + } + + // Convert to e2m1 values. + uint32_t e2m1Vec = fp32_vec_to_e2m1(fp2Vals); + + // Write the e2m1 values to global memory. + return e2m1Vec; +#else + return 0; +#endif +} + +} // namespace utils + +template +__device__ __forceinline__ vec_t vec_add(const vec_t& a, + const vec_t& b) { + vec_t ret; +#pragma unroll + for (int i = 0; i < VEC_SIZE; ++i) { + ret[i] = static_cast(a[i]) + static_cast(b[i]); + } + return ret; +} + +enum class AllReduceFusionPattern : int { + kAllReduce = 0, + kARResidualRMSNorm = 1, + kARResidualRMSNormFP8Quant = 2, + kARResidualRMSNormFP4Quant = 3, + // The difference between these two and the standard version is that the NormOut version outputs + // the result of the norm. + kARResidualRMSNormOutFP8Quant = 4, + kARResidualRMSNormOutFP4Quant = 5 +}; + +enum class QuantType : int { + kNone = 0, + kFP8 = 1, + kFP4 = 2, +}; + +template +struct FusionPatternTraits; + +#define DEFINE_FUSION_PATTERN_TRAITS(pattern, hasAllReduceOut, hasResidual, hasResidualOut, \ + hasRMSNorm, hasNormOut, quantType) \ + template <> \ + struct FusionPatternTraits { \ + static constexpr bool kHasAllReduceOut = hasAllReduceOut; \ + static constexpr bool kHasResidual = hasResidual; \ + static constexpr bool kHasResidualOut = hasResidualOut; \ + static constexpr bool kHasRMSNorm = hasRMSNorm; \ + static constexpr bool kHasNormOut = hasNormOut; \ + static constexpr QuantType kQuantType = quantType; \ + }; + +DEFINE_FUSION_PATTERN_TRAITS(AllReduceFusionPattern::kAllReduce, true, false, false, false, false, + QuantType::kNone); +DEFINE_FUSION_PATTERN_TRAITS(AllReduceFusionPattern::kARResidualRMSNorm, false, true, true, true, + true, QuantType::kNone); +DEFINE_FUSION_PATTERN_TRAITS(AllReduceFusionPattern::kARResidualRMSNormFP8Quant, false, true, true, + true, false, QuantType::kFP8); +DEFINE_FUSION_PATTERN_TRAITS(AllReduceFusionPattern::kARResidualRMSNormFP4Quant, false, true, true, + true, false, QuantType::kFP4); +DEFINE_FUSION_PATTERN_TRAITS(AllReduceFusionPattern::kARResidualRMSNormOutFP8Quant, false, true, + true, true, true, QuantType::kFP8); +DEFINE_FUSION_PATTERN_TRAITS(AllReduceFusionPattern::kARResidualRMSNormOutFP4Quant, false, true, + true, true, true, QuantType::kFP4); +#undef DEFINE_FUSION_PATTERN_TRAITS + +template +constexpr bool HasResidual = FusionPatternTraits::kHasResidual; +template +constexpr bool HasRMSNorm = FusionPatternTraits::kHasRMSNorm; +template +constexpr bool HasAllReduceOut = FusionPatternTraits::kHasAllReduceOut; +template +constexpr bool HasResidualOut = FusionPatternTraits::kHasResidualOut; +template +constexpr bool HasNormOut = FusionPatternTraits::kHasNormOut; +template +constexpr QuantType GetQuantType = FusionPatternTraits::kQuantType; + +template +struct AllReduceFusionParams { + int nranks; + int rank; + int size; + int hidden_dim; + void** workspace; + void* allreduce_in; + void* allreduce_out; + void* residual_in; + void* residual_out; + void* norm_out; + void* quant_out; + void* scale_out; + void* rms_gamma; + float rms_eps; + float* scale_factor; + bool use_oneshot; + FP4QuantizationSFLayout layout = FP4QuantizationSFLayout::SWIZZLED; + cudaStream_t stream; + AllReduceFusionPattern pattern; + bool trigger_completion_at_end = true; +}; + +template +struct SyncComm { + __device__ __forceinline__ SyncComm(void** workspace) { + counter_ptr = &reinterpret_cast(workspace[NRanks * 3])[0]; + flag_ptr = &reinterpret_cast(workspace[NRanks * 3])[1]; + flag_value = *flag_ptr; + for (int r = 0; r < NRanks; ++r) { + comm_bufs[r] = workspace[r]; + barrier_flags[r] = workspace[NRanks + r]; + } + __syncthreads(); + if (threadIdx.x == 0) { + atomicAdd(counter_ptr, 1); + } + } + + __device__ __forceinline__ void update(int new_flag_value) { + if (blockIdx.x == 0 && threadIdx.x == 0) { + while (*reinterpret_cast(counter_ptr) != gridDim.x) { + } + *flag_ptr = new_flag_value; + *counter_ptr = 0; + } + } + + int* counter_ptr; + int* flag_ptr; + void* comm_bufs[NRanks]; + void* barrier_flags[NRanks]; + int flag_value; +}; + +template +struct LamportComm { + __device__ __forceinline__ LamportComm(void** workspace, int rank) { + counter_ptr = &reinterpret_cast(workspace[NRanks * 3])[0]; + flag_ptr = &reinterpret_cast(workspace[NRanks * 3])[2]; + clear_ptr = &reinterpret_cast(workspace[NRanks * 3])[4]; + flag_value = *flag_ptr; + int comm_size = reinterpret_cast(workspace[NRanks * 3])[3]; + clear_size = *clear_ptr; + int data_offset = flag_value % 3; + int clear_offset = (flag_value + 2) % 3; + for (int r = 0; r < NRanks; ++r) { + data_bufs[r] = reinterpret_cast(workspace[2 * NRanks + r]) + + static_cast(data_offset) * comm_size; + } + clear_buf = reinterpret_cast(workspace[2 * NRanks + rank]) + clear_offset * comm_size; + __syncthreads(); + if (threadIdx.x == 0) { + atomicAdd(counter_ptr, 1); + } + } + + __device__ __forceinline__ void update(int new_clear_size) { + if (blockIdx.x == 0 && threadIdx.x == 0) { + while (*reinterpret_cast(counter_ptr) != gridDim.x) { + } + *flag_ptr = (flag_value + 1) % 3; + *clear_ptr = new_clear_size; + *counter_ptr = 0; + } + } + + int* counter_ptr; + int* flag_ptr; + int* clear_ptr; + uint8_t* data_bufs[NRanks]; + uint8_t* clear_buf; + int clear_size; + int flag_value; +}; + +template +class Barrier { + public: + __device__ __forceinline__ Barrier(int rank, SyncComm const& comm) { + if (threadIdx.x < NRanks) { + m_flag_value = comm.flag_value; + int current_rank = rank; + int target_rank = threadIdx.x; + m_target_flag = reinterpret_cast(comm.barrier_flags[target_rank]) + current_rank; + m_current_flag = reinterpret_cast(comm.barrier_flags[current_rank]) + + blockIdx.x * NRanks + target_rank; + } + } + + __device__ __forceinline__ void sync() { + __syncthreads(); + if (threadIdx.x < NRanks) { + m_flag_value = next_flag(m_flag_value); + // To avoid the ABA problem, we need to synchronize the correct flag value to all + // barrier_flags, even if the corresponding CTA has not been launched. + for (int flag_idx = blockIdx.x; flag_idx < details::kBarrierFlagCount; + flag_idx += gridDim.x) { + st_flag(m_target_flag + flag_idx * NRanks, m_flag_value); + } + while (ld_flag(m_current_flag) == prev_flag(m_flag_value)) { + } + } + __syncthreads(); + } + + protected: + __device__ __forceinline__ void st_flag(int* addr, int flag) { + asm volatile("st.global.release.sys.b32 [%1], %0;" ::"r"(flag), "l"(addr)); + } + + __device__ __forceinline__ int ld_flag(int* addr) { + int flag; + asm volatile("ld.global.acquire.sys.b32 %0, [%1];" : "=r"(flag) : "l"(addr)); + return flag; + } + + __device__ __forceinline__ int next_flag(int flag) { return flag == 2 ? 0 : flag + 1; } + + __device__ __forceinline__ int prev_flag(int flag) { return flag == 0 ? 2 : flag - 1; } + + public: + int m_flag_value; + + private: + int* m_target_flag; + int* m_current_flag; +}; + +template +class FusedOp { + static constexpr int VEC_SIZE = details::kBytesPerAccess / sizeof(T); + + public: + __device__ __forceinline__ FusedOp(AllReduceFusionParams const& params, int access_id, + int access_id_in_token) + : m_params(params), m_access_id(access_id), m_access_id_in_token(access_id_in_token) { + if constexpr (HasRMSNorm) { + m_gamma_val.load(reinterpret_cast(params.rms_gamma) + m_access_id_in_token * VEC_SIZE); + } + if constexpr (HasResidual) { + m_residual_val.load(reinterpret_cast(params.residual_in) + m_access_id * VEC_SIZE); + } + if constexpr (GetQuantType == QuantType::kFP8) { + m_scale_factor = 1.f / *(params.scale_factor); + } else if constexpr (GetQuantType == QuantType::kFP4) { + m_scale_factor = *(params.scale_factor); + } + } + + // template + __device__ __forceinline__ void update(int access_id) { + if (m_access_id != access_id) { + m_access_id = access_id; + if constexpr (HasResidual) { + m_residual_val.load(reinterpret_cast(m_params.residual_in) + m_access_id * VEC_SIZE); + } + } + } + + // template + __device__ __forceinline__ void operator()(vec_t val, int token_id) { + if constexpr (HasAllReduceOut) { + val.store(reinterpret_cast(m_params.allreduce_out) + m_access_id * VEC_SIZE); + } + if constexpr (HasResidual) { + val = vec_add(val, m_residual_val); + if constexpr (HasResidualOut) { + val.store(reinterpret_cast(m_params.residual_out) + m_access_id * VEC_SIZE); + } + } + if constexpr (HasRMSNorm) { + val = rms_norm(val, m_gamma_val); + if constexpr (HasNormOut) { + val.store(reinterpret_cast(m_params.norm_out) + m_access_id * VEC_SIZE); + } + } + + if constexpr (GetQuantType == QuantType::kFP4) { + // NOTE(Yingyi): might update later + auto sf_out = utils::cvt_quant_to_fp4_get_sf_out_offset( + std::nullopt /* batchIdx */, token_id, m_access_id_in_token, std::nullopt /* numRows */, + m_params.hidden_dim, reinterpret_cast(m_params.scale_out), m_params.layout); + reinterpret_cast(m_params.quant_out)[m_access_id] = + utils::cvt_warp_fp16_to_fp4(val, m_scale_factor, sf_out); + } else if constexpr (GetQuantType == QuantType::kFP8) { + using PackedQuantizedType = std::conditional_t, float, float2>; + PackedQuantizedType ret; +#pragma unroll + for (int i = 0; i < VEC_SIZE; ++i) { + reinterpret_cast<__nv_fp8_e4m3*>(&ret)[i] = static_cast<__nv_fp8_e4m3>( + static_cast(reinterpret_cast(&val)[i]) * m_scale_factor); + } + reinterpret_cast(m_params.quant_out)[m_access_id] = ret; + } else { + static_assert(GetQuantType == QuantType::kNone, "Invalid quant type"); + } + } + + protected: + __device__ __forceinline__ vec_t rms_norm(vec_t const& residual, + vec_t const& gamma) { + __shared__ float s_val; + vec_t norm_out; + float acc = 0.f; +#pragma unroll + for (int i = 0; i < VEC_SIZE; ++i) { + float v = static_cast(reinterpret_cast(&residual)[i]); + acc += v * v; + } + utils::blockReduceSumV2(&acc); +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + namespace cg = cooperative_groups; + cg::cluster_group cluster = cg::this_cluster(); + if (cluster.num_blocks() > 1) { + if (threadIdx.x == 0) { + s_val = acc; + acc = 0.f; + } + cluster.sync(); + if (threadIdx.x == 0) { + for (int i = 0; i < cluster.num_blocks(); ++i) { + acc += *cluster.map_shared_rank(&s_val, i); + } + } + cluster.sync(); + } +#endif + if (threadIdx.x == 0) { + s_val = rsqrtf(acc / m_params.hidden_dim + m_params.rms_eps); + } + __syncthreads(); +#pragma unroll + for (int i = 0; i < VEC_SIZE; ++i) { + reinterpret_cast(&norm_out)[i] = + static_cast(static_cast(reinterpret_cast(&residual)[i]) * s_val * + static_cast(reinterpret_cast(&gamma)[i])); + } + return norm_out; + } + + private: + AllReduceFusionParams const& m_params; + int m_access_id; + int m_access_id_in_token; + float m_scale_factor; + vec_t m_residual_val; + vec_t m_gamma_val; +}; + +template +struct neg_zero { + static constexpr T value = -T(0); +}; + +template <> +struct neg_zero { + static constexpr unsigned short neg_zero_bits = 0x8000U; + static constexpr __half value = __half_raw{neg_zero_bits}; +}; + +template <> +struct neg_zero { + static constexpr unsigned short neg_zero_bits = 0x8000U; + static constexpr __nv_bfloat16 value = __nv_bfloat16_raw{neg_zero_bits}; +}; + +template <> +struct neg_zero { + static constexpr unsigned int neg_zero_bits = 0x80000000U; + static constexpr float value = -0.0f; +}; + +template +__device__ static constexpr T neg_zero_v = neg_zero::value; + +template +__device__ bool is_negative_zero(T) { + return false; +} + +// float specialization +template <> +__device__ bool is_negative_zero(float x) { + return (__float_as_int(x) == 0x80000000); +} + +// double specialization +template <> +__device__ bool is_negative_zero(double x) { + return (__double_as_longlong(x) == 0x8000000000000000ULL); +} + +// __half specialization +template <> +__device__ bool is_negative_zero<__half>(__half x) { + return (__half_as_ushort(x) == 0x8000); +} + +// __nv_bfloat16 specialization +template <> +__device__ bool is_negative_zero<__nv_bfloat16>(__nv_bfloat16 x) { + return (__bfloat16_as_ushort(x) == 0x8000); +} + +template +__device__ __forceinline__ bool has_neg_zero(const vec_t& vec) { +#pragma unroll + for (int i = 0; i < VEC_SIZE; ++i) { + if (is_negative_zero(vec[i])) { + return true; + } + } + return false; +} + +template +__device__ __forceinline__ void remove_neg_zero(vec_t& vec) { +#pragma unroll + for (int i = 0; i < VEC_SIZE; ++i) { + vec[i] = (is_negative_zero(vec[i])) ? static_cast(0.f) : vec[i]; + } +} + +template +__device__ __forceinline__ void set_neg_zero(T* addr) { + vec_t val; + val.fill(neg_zero_v); + val.store_global_volatile(addr); +} + +template +__device__ __forceinline__ vec_t allreduce_sum(vec_t* vals) { + if constexpr (Fp32Acc) { + static_assert(!std::is_same_v); + float acc_f32[VEC_SIZE]; +#pragma unroll + for (int i = 0; i < VEC_SIZE; ++i) { + acc_f32[i] = static_cast(reinterpret_cast(&vals[0])[i]); + } +#pragma unroll + for (int r = 1; r < NRanks; ++r) { +#pragma unroll + for (int i = 0; i < VEC_SIZE; ++i) { + acc_f32[i] += static_cast(reinterpret_cast(&vals[r])[i]); + } + } + vec_t acc; +#pragma unroll + for (int i = 0; i < VEC_SIZE; ++i) { + acc[i] = static_cast(acc_f32[i]); + } + return acc; + } else { + vec_t acc = vals[0]; +#pragma unroll + for (int r = 1; r < NRanks; ++r) { + acc = vec_add(acc, vals[r]); + } + return acc; + } +} + +template +class IndexHelper { + public: + __device__ __forceinline__ IndexHelper(AllReduceFusionParams const& params) { + static constexpr int VEC_SIZE = details::kBytesPerAccess / sizeof(T); +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + namespace cg = cooperative_groups; + cg::cluster_group cluster = cg::this_cluster(); + cg::grid_group grid = cg::this_grid(); + token_id = grid.cluster_rank(); + access_id_in_token = cluster.thread_rank(); + token_stride = grid.num_clusters(); +#else + token_id = blockIdx.x; + access_id_in_token = threadIdx.x; + token_stride = gridDim.x; +#endif + access_id = token_id * params.hidden_dim / VEC_SIZE + access_id_in_token; + access_stride = token_stride * params.hidden_dim / VEC_SIZE; + tot_access = params.size / VEC_SIZE; + } + + int token_id; + int access_id_in_token; + int token_stride; + int access_id; + int access_stride; + int tot_access; +}; + +template +__global__ void allreduce_fusion_kernel_oneshot_lamport(AllReduceFusionParams params) { + static constexpr int VEC_SIZE = details::kBytesPerAccess / sizeof(T); + IndexHelper index_helper(params); + int token_id = index_helper.token_id; + int access_id_in_token = index_helper.access_id_in_token; + int token_stride = index_helper.token_stride; + int access_id = index_helper.access_id; + int access_stride = index_helper.access_stride; + int tot_access = index_helper.tot_access; + vec_t clear_vec; + clear_vec.fill(neg_zero_v); + FusedOp fused_op(params, access_id, access_id_in_token); + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + cudaGridDependencySynchronize(); + if constexpr (!TriggerCompletionAtEnd) { + cudaTriggerProgrammaticLaunchCompletion(); + } +#endif + LamportComm comm(params.workspace, params.rank); + int clear_access = comm.clear_size / VEC_SIZE; + + for (int idx = access_id; idx < tot_access; idx += access_stride) { + vec_t val; + val.load(reinterpret_cast(params.allreduce_in) + idx * VEC_SIZE); + remove_neg_zero(val); +#pragma unroll + for (int r = 0; r < NRanks; ++r) { + // Push data to other ranks + val.store(reinterpret_cast(comm.data_bufs[r]) + + (params.rank * tot_access + idx) * VEC_SIZE); + } + } + for (int idx = access_id; idx < clear_access; idx += access_stride) { + // Clear comm buffer that previous kernel used + clear_vec.store(reinterpret_cast(comm.clear_buf) + idx * VEC_SIZE); + } + + for (int idx = access_id, tidx = token_id; idx < tot_access; + idx += access_stride, tidx += token_stride) { + fused_op.update(idx); + vec_t vals[NRanks]; + bool done = false; + + while (!done) { + done = true; +#pragma unroll + for (int r = 0; r < NRanks; ++r) { + // LDG.128 from local rank + vals[r].load_global_volatile(reinterpret_cast(comm.data_bufs[params.rank]) + + (r * tot_access + idx) * VEC_SIZE); + done &= !has_neg_zero(vals[r]); + } + } + vec_t sum_val = allreduce_sum(vals); + fused_op(sum_val, tidx); + } + + comm.update(params.size * NRanks); + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + if constexpr (TriggerCompletionAtEnd) { + cudaTriggerProgrammaticLaunchCompletion(); + } +#endif +} + +template +__global__ void allreduce_fusion_kernel_twoshot_sync(AllReduceFusionParams params, + std::array begin_tokens, + std::array token_num_per_ranks) { + static constexpr int VEC_SIZE = details::kBytesPerAccess / sizeof(T); + IndexHelper index_helper(params); + int token_id = index_helper.token_id; + int access_id_in_token = index_helper.access_id_in_token; + int token_stride = index_helper.token_stride; + int access_id = index_helper.access_id; + int access_stride = index_helper.access_stride; + int tot_access = index_helper.tot_access; + FusedOp fused_op(params, access_id, access_id_in_token); +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + cudaGridDependencySynchronize(); +#endif + SyncComm comm(params.workspace); +#pragma unroll + for (int r = 0; r < NRanks; ++r) { + int comm_access_id = access_id + begin_tokens[r] * params.hidden_dim / VEC_SIZE; + int comm_tot_access = (begin_tokens[r] + token_num_per_ranks[r]) * params.hidden_dim / VEC_SIZE; + for (int idx = comm_access_id; idx < comm_tot_access; idx += access_stride) { + reinterpret_cast(comm.comm_bufs[params.rank])[idx] = + reinterpret_cast(params.allreduce_in)[idx]; + } + } + Barrier barrier(params.rank, comm); + barrier.sync(); + int comm_access_id = access_id + begin_tokens[params.rank] * params.hidden_dim / VEC_SIZE; + int comm_tot_access = + (begin_tokens[params.rank] + token_num_per_ranks[params.rank]) * params.hidden_dim / VEC_SIZE; + for (int idx = comm_access_id; idx < comm_tot_access; idx += access_stride) { + vec_t vals[NRanks]; +#pragma unroll + for (int r = 0; r < NRanks; ++r) { + vals[r].load(reinterpret_cast(comm.comm_bufs[r]) + idx * VEC_SIZE); + } + vec_t sum_val = allreduce_sum(vals); +#pragma unroll + for (int r = 0; r < NRanks; ++r) { + sum_val.store(reinterpret_cast(comm.comm_bufs[r]) + (tot_access + idx) * VEC_SIZE); + } + } + barrier.sync(); +#pragma unroll + for (int r = 0; r < NRanks; ++r) { + int comm_access_id = access_id + begin_tokens[r] * params.hidden_dim / VEC_SIZE; + int comm_token_id = token_id + begin_tokens[r]; + int comm_tot_access = (begin_tokens[r] + token_num_per_ranks[r]) * params.hidden_dim / VEC_SIZE; + for (int idx = comm_access_id, tidx = comm_token_id; idx < comm_tot_access; + idx += access_stride, tidx += token_stride) { + fused_op.update(idx); + vec_t sum_val; + sum_val.load(reinterpret_cast(comm.comm_bufs[params.rank]) + + (tot_access + idx) * VEC_SIZE); + fused_op(sum_val, tidx); + } + } + comm.update(barrier.m_flag_value); +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + cudaTriggerProgrammaticLaunchCompletion(); +#endif +} + +int get_sm_count() { + static int sm_count = 0; + if (sm_count == 0) { + int device_id; + FLASHINFER_CUDA_CALL(cudaGetDevice(&device_id)); + cudaDeviceProp device_prop; + cudaGetDeviceProperties(&device_prop, device_id); + sm_count = device_prop.multiProcessorCount; + } + return sm_count; +} + +template +cudaError_t launch_oneshot_lamport(AllReduceFusionParams const& params, + cudaLaunchConfig_t& cfg) { + FLASHINFER_CUDA_CALL(cudaLaunchKernelEx( + &cfg, + allreduce_fusion_kernel_oneshot_lamport, + params)); + return cudaSuccess; +} + +template +cudaError_t launch_twoshot_sync(AllReduceFusionParams const& params, cudaLaunchConfig_t& cfg, + std::array begin_tokens, + std::array token_num_per_ranks) { + FLASHINFER_CUDA_CALL( + cudaLaunchKernelEx(&cfg, allreduce_fusion_kernel_twoshot_sync, + params, begin_tokens, token_num_per_ranks)); + return cudaSuccess; +} + +bool use_oneshot(int token_num) { return token_num <= details::kOneShotMaxToken; } + +template +cudaError_t allreduce_fusion_kernel_launcher(AllReduceFusionParams const& params, + bool launch_with_pdl) { + static constexpr int VEC_SIZE = details::kBytesPerAccess / sizeof(T); + FLASHINFER_CHECK(params.size % params.hidden_dim == 0, "params.size % params.hidden_dim != 0"); + FLASHINFER_CHECK(params.hidden_dim % VEC_SIZE == 0, "params.hidden_dim % VEC_SIZE != 0"); + static int SM = utils::getSMVersion(); + int token_num = params.size / params.hidden_dim; + bool oneshot = params.use_oneshot; + int cluster_num = token_num; + std::array begin_tokens, token_num_per_ranks; + if (!oneshot) { + int remaining_token = token_num % NRanks; + int token_num_per_rank = token_num / NRanks; + cluster_num = token_num_per_rank; + if (remaining_token) { + cluster_num++; + } + for (int r = 0; r < NRanks; ++r) { + begin_tokens[r] = r * token_num_per_rank + (remaining_token > r ? r : remaining_token); + token_num_per_ranks[r] = token_num_per_rank + (remaining_token > r ? 1 : 0); + } + } + int threads_per_token = params.hidden_dim / VEC_SIZE; + int cluster_size; + if (SM >= 90) { + cluster_size = 8; + } else { + cluster_size = 1; + } + while (threads_per_token % cluster_size != 0 && cluster_size > 1) { + cluster_size /= 2; + } + int threads_per_block = threads_per_token / cluster_size; + while (threads_per_block < 128 && cluster_size >= 2) { + threads_per_block *= 2; + cluster_size /= 2; + } + FLASHINFER_CHECK(oneshot || threads_per_block >= params.nranks, + "not oneshot, or threads_per_block < nranks"); + int block_size = threads_per_block; + FLASHINFER_CHECK(block_size <= 1024 && cluster_size > 0, + "block_size > 1024 or cluster_size <= 0"); + int sm_count = get_sm_count(); + int grid_size = (std::min(sm_count, cluster_num * cluster_size) / cluster_size) * cluster_size; + cudaLaunchConfig_t cfg; + cudaLaunchAttribute attribute[2]; + cfg.gridDim = grid_size; + cfg.blockDim = block_size; + cfg.dynamicSmemBytes = 0; + cfg.stream = params.stream; + attribute[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attribute[0].val.programmaticStreamSerializationAllowed = launch_with_pdl ? 1 : 0; + attribute[1].id = cudaLaunchAttributeClusterDimension; + attribute[1].val.clusterDim.x = cluster_size; + attribute[1].val.clusterDim.y = 1; + attribute[1].val.clusterDim.z = 1; + cfg.attrs = attribute; + cfg.numAttrs = SM >= 90 ? 2 : 0; + if (oneshot) { + bool trigger_completion_at_end = params.trigger_completion_at_end; + if (trigger_completion_at_end) { + FLASHINFER_CUDA_CALL( + (launch_oneshot_lamport(params, cfg))); + } else { + FLASHINFER_CUDA_CALL( + (launch_oneshot_lamport(params, cfg))); + } + } else { + FLASHINFER_CUDA_CALL((launch_twoshot_sync( + params, cfg, begin_tokens, token_num_per_ranks))); + } + return cudaSuccess; +} + +template +cudaError_t allreduce_fusion_op(AllReduceFusionParams const& params, bool launch_with_pdl, + bool fp32_acc) { +#define DISPATCH_ACC_TYPE(T, Pattern, NRanks) \ + if constexpr (std::is_same_v) { \ + return allreduce_fusion_kernel_launcher(params, launch_with_pdl); \ + } else { \ + if (fp32_acc) { \ + return allreduce_fusion_kernel_launcher(params, launch_with_pdl); \ + } else { \ + return allreduce_fusion_kernel_launcher(params, launch_with_pdl); \ + } \ + } + +#define DISPATCH_PATTERN(T, NRanks) \ + switch (params.pattern) { \ + case AllReduceFusionPattern::kAllReduce: \ + DISPATCH_ACC_TYPE(T, AllReduceFusionPattern::kAllReduce, NRanks); \ + break; \ + case AllReduceFusionPattern::kARResidualRMSNorm: \ + DISPATCH_ACC_TYPE(T, AllReduceFusionPattern::kARResidualRMSNorm, NRanks); \ + break; \ + case AllReduceFusionPattern::kARResidualRMSNormFP8Quant: \ + DISPATCH_ACC_TYPE(T, AllReduceFusionPattern::kARResidualRMSNormFP8Quant, NRanks); \ + break; \ + case AllReduceFusionPattern::kARResidualRMSNormFP4Quant: \ + if constexpr (!std::is_same_v) { \ + DISPATCH_ACC_TYPE(T, AllReduceFusionPattern::kARResidualRMSNormFP4Quant, NRanks); \ + } else { \ + FLASHINFER_CHECK(false, "FP4Quant pattern cannot work with DType=float!"); \ + } \ + break; \ + case AllReduceFusionPattern::kARResidualRMSNormOutFP8Quant: \ + DISPATCH_ACC_TYPE(T, AllReduceFusionPattern::kARResidualRMSNormOutFP8Quant, NRanks); \ + break; \ + case AllReduceFusionPattern::kARResidualRMSNormOutFP4Quant: \ + if constexpr (!std::is_same_v) { \ + DISPATCH_ACC_TYPE(T, AllReduceFusionPattern::kARResidualRMSNormOutFP4Quant, NRanks); \ + } else { \ + FLASHINFER_CHECK(false, "OutFP4Quant pattern cannot work with DType=float!"); \ + } \ + break; \ + default: \ + FLASHINFER_CHECK(false, "Unsupported allreduce fusion pattern!"); \ + } + + switch (params.nranks) { + case 2: + DISPATCH_PATTERN(T, 2); + break; + case 4: + DISPATCH_PATTERN(T, 4); + break; + case 8: + DISPATCH_PATTERN(T, 8); + break; + case 16: + DISPATCH_PATTERN(T, 16); + break; + default: + FLASHINFER_ERROR( + "allreduce_fusion_kernel: unsupported ranks number! Supported ranks: 2, 4, 8, 16."); + } +} + +} // namespace trtllm_allreduce_fusion + +} // namespace flashinfer diff --git a/include/flashinfer/comm/trtllm_alltoall.cuh b/include/flashinfer/comm/trtllm_alltoall.cuh new file mode 100644 index 0000000000000000000000000000000000000000..10c400528c35f605a07509899fd9accb260f32e9 --- /dev/null +++ b/include/flashinfer/comm/trtllm_alltoall.cuh @@ -0,0 +1,925 @@ +/* + * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include + +#include + +#include "../exception.h" +#include "../logging.h" +#include "../utils.cuh" +#include "../vec_dtypes.cuh" + +namespace cg = cooperative_groups; + +namespace flashinfer { + +namespace trtllm_alltoall { + +#ifdef __CUDACC__ +#define ALIGN_256 __align__(256) +#else +#define ALIGN_256 alignas(256) +#endif + +struct ALIGN_256 MoeCommFifoConnInfo { + volatile uint64_t head; // write position + volatile uint64_t tail; // read position +}; + +constexpr int WARP_SIZE = 32; +constexpr uint32_t WARP_MASK = 0xffffffff; + +constexpr int RECV_FIFO_DEPTH = 8; +constexpr int RECV_FIFO_ENTRY_BYTES = 256 * 1024; +constexpr int RECV_FIFO_ENTRY_U64 = RECV_FIFO_ENTRY_BYTES / sizeof(uint64_t); +constexpr int RECV_FIFO_TOTAL_BYTES = RECV_FIFO_DEPTH * RECV_FIFO_ENTRY_BYTES; +constexpr int RECV_FIFO_TOTAL_U64 = RECV_FIFO_TOTAL_BYTES / sizeof(uint64_t); + +static int getMultiProcessorCount() { + int device_id; + int multi_processor_count; + FLASHINFER_CUDA_CALL(cudaGetDevice(&device_id)); + FLASHINFER_CUDA_CALL( + cudaDeviceGetAttribute(&multi_processor_count, cudaDevAttrMultiProcessorCount, device_id)); + return multi_processor_count; +} + +class AllToAllChannelCommunicatorBase { + public: + static constexpr int GROUP_COUNT_PER_BLOCK = 8; + static_assert(GROUP_COUNT_PER_BLOCK <= 8, + "GROUP_COUNT_PER_BLOCK must be less than or equal to 8"); + static constexpr int WARP_PER_GROUP = 2; + static constexpr int U64_DATA_REG_PER_THREAD = 8; + // A packet is a warp-sized chunk of data that is sent or received in one go, + // but may be split into multiple 64-bit registers, the number of which is + // U64_DATA_REG_PER_THREAD. + static constexpr int PACKET_SIZE_IN_U64 = WARP_SIZE * U64_DATA_REG_PER_THREAD; + static constexpr int PACKET_SIZE_IN_BYTES = PACKET_SIZE_IN_U64 * sizeof(uint64_t); + static constexpr int DATA_PAYLOAD_SIZE_PER_PACKET_IN_U64 = + (WARP_SIZE - 2) * U64_DATA_REG_PER_THREAD; + static constexpr int DATA_PAYLOAD_SIZE_PER_PACKET = + DATA_PAYLOAD_SIZE_PER_PACKET_IN_U64 * sizeof(uint64_t); + static constexpr int U64_ELT_COUNT_PER_PACKET = PACKET_SIZE_IN_BYTES / sizeof(uint64_t); + + static constexpr int PACKET_COUNT_PER_FIFO_ENTRY = RECV_FIFO_ENTRY_BYTES / PACKET_SIZE_IN_BYTES; + + static constexpr int GROUP_MAX_INDICE_COUNT = + RECV_FIFO_ENTRY_BYTES / sizeof(uint64_t) / (WARP_SIZE * U64_DATA_REG_PER_THREAD); + + struct GroupSharedBuffer { + int groupIndiceBuffer[GROUP_MAX_INDICE_COUNT]; + int groupStartIndice; + int groupEndIndice; + }; + + static void setMaxUsableSmCount(int maxUsableSmCount) { + FLASHINFER_CHECK(AllToAllChannelCommunicatorBase::maxSmCountUsed == false, + "setMaxUsableSmCount can be called only before it is used"); + int smCount = getMultiProcessorCount(); + if (maxUsableSmCount > smCount) { + FLASHINFER_LOG_WARN( + "setMaxUsableSmCount, maxUsableSmCount=%d, larger than smCount=%d, using smCount instead", + maxUsableSmCount, smCount); + maxUsableSmCount = smCount; + } + AllToAllChannelCommunicatorBase::maxSmCount = maxUsableSmCount; + } + + static int getMaxUsableSmCount() { + AllToAllChannelCommunicatorBase::maxSmCountUsed = true; + if (AllToAllChannelCommunicatorBase::maxSmCount == -1) { + int smCount = getMultiProcessorCount(); + AllToAllChannelCommunicatorBase::maxSmCount = smCount; + } + return AllToAllChannelCommunicatorBase::maxSmCount; + } + + static int computeMoeCommChannelCount(int epSize) { + int smCount = getMaxUsableSmCount(); + int blockCountPerChannel = (epSize + GROUP_COUNT_PER_BLOCK - 1) / GROUP_COUNT_PER_BLOCK; + blockCountPerChannel *= 2; // for send and recv + FLASHINFER_CHECK(blockCountPerChannel <= smCount, + "GPU should support at lease one channel, usableSmCount=%d", smCount); + int perferredChannel = smCount / 2 / blockCountPerChannel; // use half SMs for communication + int channelCount = std::max(perferredChannel, 1); // at lease one channel + return channelCount; + } + + static int getMoeCommChannelCount(int epSize) { + static std::map channelCountMap{}; + auto iter = channelCountMap.find(epSize); + if (iter == channelCountMap.end()) { + auto channelCount = AllToAllChannelCommunicatorBase::computeMoeCommChannelCount(epSize); + channelCountMap[epSize] = channelCount; + return channelCount; + } + return iter->second; + } + + static dim3 getLaunchBlockDim() { + return dim3(WARP_SIZE * WARP_PER_GROUP, GROUP_COUNT_PER_BLOCK); + } + + static dim3 getLaunchGridDim(int epSize) { + int channelCount = AllToAllChannelCommunicatorBase::getMoeCommChannelCount(epSize); + return dim3((epSize + GROUP_COUNT_PER_BLOCK - 1) / GROUP_COUNT_PER_BLOCK, channelCount, 2); + } + + protected: + static int maxSmCount; + static bool maxSmCountUsed; +}; + +inline size_t getMoeCommWorkspaceSize(int epSize) { + int channelCount = AllToAllChannelCommunicatorBase::getMoeCommChannelCount(epSize); + return RECV_FIFO_TOTAL_BYTES * epSize * channelCount + + sizeof(MoeCommFifoConnInfo) * epSize * channelCount; +} + +struct MoeEpWorldInfo { + int epSize; + int epRank; +}; + +struct MoeExpertParallelInfo { + int expertCount = -1; + int topK = 1; +}; + +struct SendRecvDataInfo { + int vectorSizeInU64; + // pre-computed at host side for GPU kernel + int dataPacketCountPerVector; + int vectorCountPerFifoEntry; + + void ComputeDataPacketCountPerVector() { + dataPacketCountPerVector = (vectorSizeInU64 * sizeof(uint64_t) + + AllToAllChannelCommunicatorBase::DATA_PAYLOAD_SIZE_PER_PACKET - 1) / + AllToAllChannelCommunicatorBase::DATA_PAYLOAD_SIZE_PER_PACKET; + } + + void ComputeVectorCountPerFifoEntry() { + ComputeDataPacketCountPerVector(); + vectorCountPerFifoEntry = + AllToAllChannelCommunicatorBase::PACKET_COUNT_PER_FIFO_ENTRY / dataPacketCountPerVector; + } + + void DoPreCompute() { + ComputeDataPacketCountPerVector(); + ComputeVectorCountPerFifoEntry(); + assert(vectorCountPerFifoEntry <= AllToAllChannelCommunicatorBase::GROUP_MAX_INDICE_COUNT); + } +}; + +// struct holding Send/Recv data pointer and its displacement information. +struct SendRecvDispls { + uint64_t* dataPtr; + int const* rankCountCumSum; // length = epSize + int const* rankLocalIndices; // length = rankCountCumSum[epRank] - rankCountCumSum[epRank - 1] if + // epRank > 0 else rankCountCumSum[epRank] + int vectorStrideInU64; + +#ifdef __CUDACC__ + __inline__ __device__ int getCount(int rank) const { + return rank == 0 ? rankCountCumSum[rank] : rankCountCumSum[rank] - rankCountCumSum[rank - 1]; + } + + __inline__ __device__ int getRankStart(int rank) const { + return rank == 0 ? 0 : rankCountCumSum[rank - 1]; + } + + __inline__ __device__ int getRealVectorIndice(int globalVectorIndex) const { + return rankLocalIndices[globalVectorIndex]; + } + + __inline__ __device__ uint64_t* getVectorDataPtr(int realVectorIndex) const { + return dataPtr + realVectorIndex * vectorStrideInU64; + } +#endif +}; + +struct MoeCommWorkspace { + uint64_t* workspacePtr; + size_t rankStrideInU64; +#ifdef __CUDACC__ + __inline__ __device__ uint64_t* getFifoBasePtr(bool isSender, int epRank, int peerRank, + int channel, int channelCount) const { + // fifo itself is in receiver's side. + if (isSender) { + return workspacePtr + peerRank * rankStrideInU64 + + (epRank * channelCount + channel) * RECV_FIFO_TOTAL_U64; + } else { + return workspacePtr + epRank * rankStrideInU64 + + (peerRank * channelCount + channel) * RECV_FIFO_TOTAL_U64; + } + } + + __inline__ __device__ MoeCommFifoConnInfo* getFifoConnInfo(bool isSender, int epRank, + int peerRank, int channel, int epSize, + int channelCount) const { + // fifoInfo is in sender's side. + uint64_t* fifoInfoPtrU64 = workspacePtr + RECV_FIFO_TOTAL_U64 * channelCount * epSize; + int strideIndice = isSender ? epRank : peerRank; + int fifoInfoIndice = isSender ? peerRank : epRank; + fifoInfoPtrU64 += strideIndice * rankStrideInU64; + MoeCommFifoConnInfo* fifoInfoPtr = (MoeCommFifoConnInfo*)fifoInfoPtrU64; + return fifoInfoPtr + fifoInfoIndice * channelCount + channel; + } +#endif +}; + +__device__ inline void barrier_sync(int name, int nThreads) { + asm volatile("barrier.sync.aligned %0, %1;" ::"r"(name), "r"(nThreads) : "memory"); +} + +inline __device__ void load128(uint64_t const* ptr, uint64_t& v0, uint64_t& v1) { + asm volatile("ld.volatile.global.v2.u64 {%0,%1}, [%2];" + : "=l"(v0), "=l"(v1) + : "l"(ptr) + : "memory"); +} + +inline __device__ void store128(uint64_t* ptr, uint64_t v0, uint64_t v1) { + asm volatile("st.volatile.global.v2.u64 [%2], {%0,%1};" ::"l"(v0), "l"(v1), "l"(ptr) : "memory"); +} + +template +class AllToAllChannelCommunicator : public AllToAllChannelCommunicatorBase { + private: + int const tid; // thread index in primitives group + int const nthreads; // number of threads in primitives group + int const wid; // lane index in warp + int const warp; // warp index in primitives group + const MoeEpWorldInfo worldInfo; + const MoeCommWorkspace workspace; + const SendRecvDataInfo sendRecvDataInfo; + const SendRecvDispls dataDispls; + int peerRank; // peer rank index + bool const flagThread; + int const group; // primitives group index + int const channel; // channel index + int const channelCount; // count of channels + + MoeCommFifoConnInfo* fifoConnInfoPtr; + uint64_t* fifoBasePtr; // pointer to fifo base address + uint64_t step; + uint64_t tailStepCache; + uint64_t regs[U64_DATA_REG_PER_THREAD]; + GroupSharedBuffer* groupSharedBuffer; + + int groupStartIndice; + int groupEndIndice; + + int sliceStartIndice; + int sliceEndIndice; + + uint64_t* stepFifoEntryPtr; + + public: + __inline__ __device__ uint64_t getFlag() { return step + 1; } + + __inline__ __device__ AllToAllChannelCommunicator(MoeEpWorldInfo const& worldInfo, + MoeCommWorkspace workspace, + SendRecvDataInfo sendRecvDataInfo, + SendRecvDispls dataDispls, + GroupSharedBuffer* groupSharedBuffer, + int channelCount) + : worldInfo(worldInfo), + nthreads(blockDim.x), + tid(threadIdx.x), + workspace(workspace), + sendRecvDataInfo(sendRecvDataInfo), + dataDispls(dataDispls), + wid(threadIdx.x % WARP_SIZE), + warp(threadIdx.x / WARP_SIZE), + peerRank(blockIdx.x * GROUP_COUNT_PER_BLOCK + threadIdx.y), + group(threadIdx.y), + channel(blockIdx.y), + flagThread(threadIdx.x % 8 == 7), + fifoConnInfoPtr(nullptr), + fifoBasePtr(nullptr), + step(0), + tailStepCache(0), + groupSharedBuffer(groupSharedBuffer), + channelCount(channelCount) {} + + __inline__ __device__ void init() { + fifoBasePtr = + workspace.getFifoBasePtr(isSender, worldInfo.epRank, peerRank, channel, channelCount); + fifoConnInfoPtr = workspace.getFifoConnInfo(isSender, worldInfo.epRank, peerRank, channel, + worldInfo.epSize, channelCount); + step = isSender ? fifoConnInfoPtr->head : fifoConnInfoPtr->tail; + tailStepCache = isSender ? fifoConnInfoPtr->tail : 0; + } + + __inline__ __device__ void computeGroupTransferRange() { + if (tid == 0) { + int rankCount = dataDispls.getCount(peerRank); + int rankStart = dataDispls.getRankStart(peerRank); + int countPerChannel = (rankCount + channelCount - 1) / channelCount; + int groupEnd = min(rankStart + (channel + 1) * countPerChannel, rankStart + rankCount); + int groupStart = min(rankStart + channel * countPerChannel, rankStart + rankCount); + groupSharedBuffer->groupStartIndice = groupStart; + groupSharedBuffer->groupEndIndice = groupEnd; + } + barrier(); + groupStartIndice = groupSharedBuffer->groupStartIndice; + groupEndIndice = groupSharedBuffer->groupEndIndice; + } + + __inline__ __device__ void loadTransferIndices() { + sliceStartIndice = groupStartIndice; + sliceEndIndice = + min(groupStartIndice + sendRecvDataInfo.vectorCountPerFifoEntry, groupEndIndice); + for (int i = groupStartIndice + tid; i < sliceEndIndice; i += WARP_SIZE * WARP_PER_GROUP) { + groupSharedBuffer->groupIndiceBuffer[i - groupStartIndice] = + dataDispls.getRealVectorIndice(i); + } + groupStartIndice = sliceEndIndice; + barrier(); + } + + __inline__ __device__ void computeSlicePtr() { + stepFifoEntryPtr = fifoBasePtr + RECV_FIFO_ENTRY_U64 * (step % RECV_FIFO_DEPTH); + } + + __inline__ __device__ void sendSlice() { + waitSend(); + int EltPer16B = 2; + int eltN = sendRecvDataInfo.vectorSizeInU64; + for (int vecId = warp + sliceStartIndice; vecId < sliceEndIndice; vecId += WARP_PER_GROUP) { + int idxInSlice = vecId - sliceStartIndice; + int vecRealIdx = groupSharedBuffer->groupIndiceBuffer[idxInSlice]; + uint64_t* src = dataDispls.getVectorDataPtr(vecRealIdx); + uint64_t* slicePtr = + stepFifoEntryPtr + + idxInSlice * sendRecvDataInfo.dataPacketCountPerVector * PACKET_SIZE_IN_U64 + 2 * wid; + for (int packetId = 0; packetId < sendRecvDataInfo.dataPacketCountPerVector; packetId++) { + int vecOff = packetId * DATA_PAYLOAD_SIZE_PER_PACKET_IN_U64; +#pragma unroll + for (int g = 0; g < U64_DATA_REG_PER_THREAD / 2; g++) { + int ix = g * WARP_SIZE - 4 * (g / 2) + wid - (g % 2) * (wid / 8); + __syncwarp(); + if (!flagThread || g % 2 == 0) { + if (ix * EltPer16B + vecOff < eltN) { + load128((uint64_t*)(src + ix * EltPer16B + vecOff), regs[2 * g + 0], regs[2 * g + 1]); + } + } + __syncwarp(); + } +#pragma unroll + for (int g = 1; g < U64_DATA_REG_PER_THREAD / 2; g += 2) { + if (flagThread) regs[2 * g] = regs[2 * g - 1]; + } + + uint64_t flag = getFlag(); + uint64_t* packetPtr = slicePtr + packetId * PACKET_SIZE_IN_U64; + __syncwarp(); +#pragma unroll + for (int u = 0; u < U64_DATA_REG_PER_THREAD; u += 2) { + store128(packetPtr + u * WARP_SIZE, regs[u], flagThread ? flag : regs[u + 1]); + } + } + } + updateSend(); + } + + __inline__ __device__ void recvSlice() { + // receiver don't need to wait since we have flag. + int EltPer16B = 2; + int eltN = sendRecvDataInfo.vectorSizeInU64; + for (int vecId = warp + sliceStartIndice; vecId < sliceEndIndice; vecId += WARP_PER_GROUP) { + int idxInSlice = vecId - sliceStartIndice; + int vecRealIdx = groupSharedBuffer->groupIndiceBuffer[idxInSlice]; + + uint64_t* dst = dataDispls.getVectorDataPtr(vecRealIdx); + uint64_t* slicePtr = + stepFifoEntryPtr + + idxInSlice * sendRecvDataInfo.dataPacketCountPerVector * PACKET_SIZE_IN_U64 + 2 * wid; + for (int packetId = 0; packetId < sendRecvDataInfo.dataPacketCountPerVector; packetId++) { + uint64_t* packetPtr = slicePtr + packetId * PACKET_SIZE_IN_U64; + int vecOff = packetId * DATA_PAYLOAD_SIZE_PER_PACKET_IN_U64; + + bool needReload; + uint64_t flag = getFlag(); + __syncwarp(); + do { + needReload = false; +#pragma unroll + for (int u = 0; u < U64_DATA_REG_PER_THREAD; u += 2) { + load128(packetPtr + u * WARP_SIZE, regs[u], regs[u + 1]); + needReload |= flagThread && (regs[u + 1] != flag); + } + } while (__any_sync(WARP_MASK, needReload)); +#pragma unroll + for (int g = 1; g < U64_DATA_REG_PER_THREAD / 2; g += 2) { + if (flagThread) regs[2 * g - 1] = regs[2 * g]; + } + +#pragma unroll + for (int g = 0; g < U64_DATA_REG_PER_THREAD / 2; g++) { + int ix = g * WARP_SIZE - 4 * (g / 2) + wid - (g % 2) * (wid / 8); + __syncwarp(); + if (!flagThread || g % 2 == 0) { + if (ix * EltPer16B + vecOff < eltN) { + store128((uint64_t*)(dst + ix * EltPer16B + vecOff), regs[2 * g + 0], + regs[2 * g + 1]); + } + } + __syncwarp(); + } + } + } + updateRecv(); + } + + __inline__ __device__ void run() { + if (peerRank >= worldInfo.epSize) { + return; + } + init(); + computeGroupTransferRange(); + while (groupStartIndice < groupEndIndice) { + loadTransferIndices(); + computeSlicePtr(); + if (isSender) { + sendSlice(); + } else { + recvSlice(); + } + } + } + + __inline__ __device__ ~AllToAllChannelCommunicator() {} + + __inline__ __device__ void barrier() { barrier_sync(15 - group, nthreads); } + + __inline__ __device__ void waitSend() { + barrier(); + while (tailStepCache + RECV_FIFO_DEPTH < step + 1) { + tailStepCache = fifoConnInfoPtr->tail; + } + barrier(); + } + + __inline__ __device__ void updateSend() { + barrier(); + if (tid == 0) { + atomicAdd_system((unsigned long long*)&fifoConnInfoPtr->head, 1); + } + barrier(); + step++; + } + + __inline__ __device__ void updateRecv() { + barrier(); + if (tid == 0) { + atomicAdd_system((unsigned long long*)&fifoConnInfoPtr->tail, 1); + } + barrier(); + step++; + } +}; + +__global__ void moeAllToAllKernel(MoeEpWorldInfo worldInfo, MoeCommWorkspace workspace, + SendRecvDataInfo sendRecvDataInfo, SendRecvDispls sendDispls, + SendRecvDispls recvDispls) { + __shared__ AllToAllChannelCommunicatorBase::GroupSharedBuffer + allGroupSharedBuffer[AllToAllChannelCommunicatorBase::GROUP_COUNT_PER_BLOCK]; + bool isSender = blockIdx.z == 0; + int channelCount = gridDim.y; + int group = threadIdx.y; + SendRecvDispls dataDispls = isSender ? sendDispls : recvDispls; + AllToAllChannelCommunicatorBase::GroupSharedBuffer* groupSharedBuffer = + &allGroupSharedBuffer[group]; + if (isSender) { + AllToAllChannelCommunicator comm(worldInfo, workspace, sendRecvDataInfo, dataDispls, + groupSharedBuffer, channelCount); + comm.run(); + } else { + AllToAllChannelCommunicator comm(worldInfo, workspace, sendRecvDataInfo, dataDispls, + groupSharedBuffer, channelCount); + comm.run(); + } +} + +void moeAllToAll(MoeEpWorldInfo worldInfo, SendRecvDataInfo sendRecvDataInfo, + SendRecvDispls sendDispls, SendRecvDispls recvDispls, MoeCommWorkspace workspace, + cudaStream_t stream) { + sendRecvDataInfo.DoPreCompute(); + FLASHINFER_CHECK(reinterpret_cast(sendDispls.dataPtr) % 16 == 0, + "sendDispls.dataPtr must be 16-byte aligned"); + FLASHINFER_CHECK(reinterpret_cast(recvDispls.dataPtr) % 16 == 0, + "recvDispls.dataPtr must be 16-byte aligned"); + dim3 block = AllToAllChannelCommunicatorBase::getLaunchBlockDim(); + dim3 grid = AllToAllChannelCommunicatorBase::getLaunchGridDim(worldInfo.epSize); + moeAllToAllKernel<<>>(worldInfo, workspace, sendRecvDataInfo, sendDispls, + recvDispls); +} + +template +__inline__ __device__ void computeSendRecvRankCountDevice( + MoeEpWorldInfo worldInfo, MoeExpertParallelInfo expertParallelInfo, int maxTokenCountPerRank, + int const* realRankTokenCountCumSum, int const* gatheredTargetRankIds, + int* sharedSendRecvRankCount, int* sendRecvRankCount) { + cg::thread_block_tile tile = + cg::tiled_partition(cg::this_thread_block()); + int laneInTile = tile.thread_rank(); + int tileId = threadIdx.x / kThreadsGroupSize; + int tileCountPerBlock = blockDim.x / kThreadsGroupSize; + + int topK = expertParallelInfo.topK; + int epRank = worldInfo.epRank; + int epSize = worldInfo.epSize; + + if (threadIdx.x == 0) { + *sharedSendRecvRankCount = 0; + } + + __syncthreads(); + int readRank = isSend ? epRank : blockIdx.x; + int compareRankId = isSend ? blockIdx.x : epRank; + int const* readRankTargetRankIds = gatheredTargetRankIds + readRank * maxTokenCountPerRank * topK; + int readRankTokenCount = maxTokenCountPerRank; + if (realRankTokenCountCumSum != nullptr) { + int readRankStart = readRank == 0 ? 0 : realRankTokenCountCumSum[readRank - 1]; + readRankTargetRankIds = gatheredTargetRankIds + readRankStart * topK; + readRankTokenCount = realRankTokenCountCumSum[readRank] - readRankStart; + } + + for (int i = tileId + blockIdx.z * tileCountPerBlock; i < readRankTokenCount; + i += tileCountPerBlock * gridDim.z) { + int targetRankId = laneInTile < topK ? readRankTargetRankIds[i * topK + laneInTile] : epSize; + bool rankMatched = (targetRankId == compareRankId); + bool hasRankMatched = tile.any(rankMatched); + if (hasRankMatched && laneInTile == 0) { + atomicAdd_block(sharedSendRecvRankCount, 1); + } + tile.sync(); + } + __syncthreads(); + if (threadIdx.x == 0) { + atomicAdd_system(sendRecvRankCount + blockIdx.x, *sharedSendRecvRankCount); + } +} + +template +__global__ void computeSendRecvRankCountKernel(MoeEpWorldInfo worldInfo, + MoeExpertParallelInfo expertParallelInfo, + int maxTokenCountPerRank, + int const* realRankTokenCountCumSum, + int const* gatheredTargetRankIds, int* sendRankCount, + int* recvRankCount) { + static_assert(kThreadsGroupSize == 1 || kThreadsGroupSize == 2 || kThreadsGroupSize == 4 || + kThreadsGroupSize == 8 || kThreadsGroupSize == 16 || kThreadsGroupSize == 32, + "Only 1, 2, 4, 8, 16, 32 threads group size supported now."); + __shared__ int sharedSendRecvRankCount; + if (blockIdx.y == 0) { + // compute send rank count + computeSendRecvRankCountDevice( + worldInfo, expertParallelInfo, maxTokenCountPerRank, realRankTokenCountCumSum, + gatheredTargetRankIds, &sharedSendRecvRankCount, sendRankCount); + } else { + // compute recv rank count + computeSendRecvRankCountDevice( + worldInfo, expertParallelInfo, maxTokenCountPerRank, realRankTokenCountCumSum, + gatheredTargetRankIds, &sharedSendRecvRankCount, recvRankCount); + } +} + +void computeSendRecvRankCount(const MoeEpWorldInfo& worldInfo, + MoeExpertParallelInfo expertParallelInfo, int maxTokenCountPerRank, + int const* realRankTokenCountCumSum, int const* gatheredTargetRankIds, + int* sendRankCount, int* recvRankCount, cudaStream_t stream) { + FLASHINFER_CHECK(expertParallelInfo.topK <= 32, + "Only topK less than or equal to 32 supported now."); + int threadsPerBlock = 1024; + auto* kernelPtr = computeSendRecvRankCountKernel<32>; + if (expertParallelInfo.topK <= 1) { + kernelPtr = computeSendRecvRankCountKernel<1>; + } else if (expertParallelInfo.topK <= 2) { + kernelPtr = computeSendRecvRankCountKernel<2>; + } else if (expertParallelInfo.topK <= 4) { + kernelPtr = computeSendRecvRankCountKernel<4>; + } else if (expertParallelInfo.topK <= 8) { + kernelPtr = computeSendRecvRankCountKernel<8>; + } else if (expertParallelInfo.topK <= 16) { + kernelPtr = computeSendRecvRankCountKernel<16>; + } + dim3 block(worldInfo.epSize, 2, 1); + kernelPtr<<>>( + worldInfo, expertParallelInfo, maxTokenCountPerRank, realRankTokenCountCumSum, + gatheredTargetRankIds, sendRankCount, recvRankCount); +} + +template +__global__ void inplaceSendRecvRankCumSumKernel(MoeEpWorldInfo worldInfo, int* sendRankCount, + int* recvRankCount) { + int* inputOutputPtr = blockIdx.x == 0 ? sendRankCount : recvRankCount; + typedef cub::BlockScan BlockScan; + __shared__ typename BlockScan::TempStorage temp_storage; + + int tid = threadIdx.x; + int threadData = tid < worldInfo.epSize ? inputOutputPtr[tid] : 0; + + BlockScan(temp_storage).InclusiveSum(threadData, threadData); + if (tid < worldInfo.epSize) { + inputOutputPtr[tid] = threadData; + } +} + +void inplaceSendRecvRankCumSum(const MoeEpWorldInfo& worldInfo, int* sendRankCount, + int* recvRankCount, cudaStream_t stream) { + int epSize = worldInfo.epSize; + int epRank = worldInfo.epRank; + + FLASHINFER_CHECK(epSize <= 1024, "Only epSize less than or equal to 1024 supported now."); + + auto* kernelPtr = inplaceSendRecvRankCumSumKernel<1024>; + int blockSize = 1024; + if (epSize <= 32) { + kernelPtr = inplaceSendRecvRankCumSumKernel<32>; + blockSize = 32; + } else if (epSize <= 64) { + kernelPtr = inplaceSendRecvRankCumSumKernel<64>; + blockSize = 64; + } else if (epSize <= 128) { + kernelPtr = inplaceSendRecvRankCumSumKernel<128>; + blockSize = 128; + } else if (epSize <= 256) { + kernelPtr = inplaceSendRecvRankCumSumKernel<256>; + blockSize = 256; + } else if (epSize <= 512) { + kernelPtr = inplaceSendRecvRankCumSumKernel<512>; + blockSize = 512; + } + kernelPtr<<<2, blockSize, 0, stream>>>(worldInfo, sendRankCount, recvRankCount); +} + +template +__inline__ __device__ void computeSendRecvIndicesDevice( + MoeEpWorldInfo worldInfo, MoeExpertParallelInfo expertParallelInfo, int maxTokenCountPerRank, + int const* realRankTokenCountCumSum, int const* gatheredTargetRankIds, + int const* sendRecvCumSum, + int* sendRecvIndices, // send or receive + int* localGatherIndices, // receive only + int* backwardRecvRankLocalIndices, // send only + int* sharedSendRecvRankStart, + typename cub::BlockScan::TempStorage& tempStorage) { + cg::thread_block_tile tile = + cg::tiled_partition(cg::this_thread_block()); + int laneInTile = tile.thread_rank(); + int tileId = threadIdx.x / kThreadsGroupSize; + int tileCountPerBlock = blockDim.x / kThreadsGroupSize; + + int topK = expertParallelInfo.topK; + int epRank = worldInfo.epRank; + int epSize = worldInfo.epSize; + + if (threadIdx.x == 0) { + *sharedSendRecvRankStart = blockIdx.x == 0 ? 0 : sendRecvCumSum[blockIdx.x - 1]; + } + + __syncthreads(); + int readRank = isSend ? epRank : blockIdx.x; + int compareRankId = isSend ? blockIdx.x : epRank; + int readRankStart = readRank * maxTokenCountPerRank; + int const* readRankTargetRankIds = gatheredTargetRankIds + readRankStart * topK; + int readRankTokenCount = maxTokenCountPerRank; + if (realRankTokenCountCumSum != nullptr) { + readRankStart = readRank == 0 ? 0 : realRankTokenCountCumSum[readRank - 1]; + readRankTargetRankIds = gatheredTargetRankIds + readRankStart * topK; + readRankTokenCount = realRankTokenCountCumSum[readRank] - readRankStart; + } + + for (int blockStartId = blockIdx.z * tileCountPerBlock; blockStartId < readRankTokenCount; + blockStartId += tileCountPerBlock * gridDim.z) { + int stepStartIndice = *sharedSendRecvRankStart; + int i = blockStartId + tileId; + int targetRankId = (laneInTile < topK && i < readRankTokenCount) + ? readRankTargetRankIds[i * topK + laneInTile] + : epSize; + bool rankMatched = (targetRankId == compareRankId); + bool hasRankMatched = tile.any(rankMatched); + unsigned int laneMask = tile.ballot(rankMatched); + int lowestLane = __ffs(laneMask) - 1; + int isMatchedLane = (hasRankMatched && laneInTile == lowestLane) ? 1 : 0; + int indice; + typedef cub::BlockScan BlockScan; + BlockScan(tempStorage).ExclusiveSum(isMatchedLane, indice); + indice += stepStartIndice; + __syncthreads(); + + if (isMatchedLane == 1) { + atomicAdd_block(sharedSendRecvRankStart, 1); + if (isSend) { + sendRecvIndices[indice] = i; + backwardRecvRankLocalIndices[indice] = i * topK + lowestLane; + } else { + sendRecvIndices[indice] = indice; + localGatherIndices[indice] = readRankStart + i; + } + } + __syncthreads(); + } +} + +template +__global__ void computeSendRecvIndicesKernel( + MoeEpWorldInfo worldInfo, MoeExpertParallelInfo expertParallelInfo, int maxTokenCountPerRank, + int const* realRankTokenCountCumSum, int const* gatheredTargetRankIds, + int const* sendRankCountCumSum, int const* recvRankCountCumSum, int* localGatherIndices, + int* sendRankLocalIndices, int* recvRankLocalIndices, int* backwardRecvRankLocalIndices) { + static_assert(kThreadsGroupSize == 1 || kThreadsGroupSize == 2 || kThreadsGroupSize == 4 || + kThreadsGroupSize == 8 || kThreadsGroupSize == 16 || kThreadsGroupSize == 32, + "Only 1, 2, 4, 8, 16, 32 threads group size supported now."); + __shared__ int sharedSendRecvRankStart; + __shared__ typename cub::BlockScan::TempStorage tempStorage; + if (blockIdx.y == 0) { + // compute send rank count + computeSendRecvIndicesDevice( + worldInfo, expertParallelInfo, maxTokenCountPerRank, realRankTokenCountCumSum, + gatheredTargetRankIds, sendRankCountCumSum, sendRankLocalIndices, localGatherIndices, + backwardRecvRankLocalIndices, &sharedSendRecvRankStart, tempStorage); + } else { + // compute recv rank count + computeSendRecvIndicesDevice( + worldInfo, expertParallelInfo, maxTokenCountPerRank, realRankTokenCountCumSum, + gatheredTargetRankIds, recvRankCountCumSum, recvRankLocalIndices, localGatherIndices, + backwardRecvRankLocalIndices, &sharedSendRecvRankStart, tempStorage); + } +} + +void computeSendRecvIndices(const MoeEpWorldInfo& worldInfo, + MoeExpertParallelInfo expertParallelInfo, int maxTokenCountPerRank, + int const* realRankTokenCountCumSum, int const* gatheredTargetRankIds, + int const* sendRankCountCumSum, int const* recvRankCountCumSum, + int* localGatherIndices, int* sendRankLocalIndices, + int* recvRankLocalIndices, int* backwardRecvRankLocalIndices, + cudaStream_t stream) { + FLASHINFER_CHECK(expertParallelInfo.topK <= 32, + "Only topK less than or equal to 32 supported now."); + int threadsPerBlock = 1024; + auto* kernelPtr = computeSendRecvIndicesKernel<32, 1024>; + if (expertParallelInfo.topK <= 1) { + kernelPtr = computeSendRecvIndicesKernel<1, 1024>; + } else if (expertParallelInfo.topK <= 2) { + kernelPtr = computeSendRecvIndicesKernel<2, 1024>; + } else if (expertParallelInfo.topK <= 4) { + kernelPtr = computeSendRecvIndicesKernel<4, 1024>; + } else if (expertParallelInfo.topK <= 8) { + kernelPtr = computeSendRecvIndicesKernel<8, 1024>; + } else if (expertParallelInfo.topK <= 16) { + kernelPtr = computeSendRecvIndicesKernel<16, 1024>; + } else if (expertParallelInfo.topK <= 32) { + kernelPtr = computeSendRecvIndicesKernel<32, 1024>; + } + dim3 block(worldInfo.epSize, 2, 1); + kernelPtr<<>>( + worldInfo, expertParallelInfo, maxTokenCountPerRank, realRankTokenCountCumSum, + gatheredTargetRankIds, sendRankCountCumSum, recvRankCountCumSum, localGatherIndices, + sendRankLocalIndices, recvRankLocalIndices, backwardRecvRankLocalIndices); +} + +cudaError_t moeAllToAllPrepareIndices( + MoeEpWorldInfo worldInfo, MoeExpertParallelInfo expertParallelInfo, int maxTokenCountPerRank, + int const* gatheredTargetRankIds, int const* realRankTokenCountCumSum, + // indices of gatheredTargetRankIds that has the local rank in topK + int* localGatherIndices, // max length = maxTokenCountPerRank * worldInfo.epSize when all ranks + // send to current rank + int* sendRankCountCumSum, // max length = worldInfo.epSize + int* sendRankLocalIndices, // max length = maxTokenCountPerRank * + // expertParallelInfo.expertCount when current rank has + // maxTokenCountPerRank tokens to send and all has expertCount dest + int* recvRankCountCumSum, // max length = worldInfo.epSize + int* recvRankLocalIndices, // max length = maxTokenCountPerRank * worldInfo.epSize when all + // ranks send to current rank + // the rankCountCumSum of combineRecv should be the same as sendRankCountCumSum + int* backwardRecvRankLocalIndices, // max length = maxTokenCountPerRank * + // expertParallelInfo.expertCount when current rank has + // maxTokenCountPerRank tokens to send and all has + // expertCount dest + cudaStream_t stream) { + FLASHINFER_CHECK(worldInfo.epSize <= 1024, + "Only worldInfo.epSize less than or equal to 1024 supported now."); + + FLASHINFER_CUDA_CALL( + cudaMemsetAsync(sendRankCountCumSum, 0, sizeof(int) * worldInfo.epSize, stream)); + FLASHINFER_CUDA_CALL( + cudaMemsetAsync(recvRankCountCumSum, 0, sizeof(int) * worldInfo.epSize, stream)); + int maxSendRanksPerToken = std::max(worldInfo.epSize, expertParallelInfo.topK); + + FLASHINFER_CUDA_CALL(cudaMemsetAsync( + localGatherIndices, -1, maxTokenCountPerRank * worldInfo.epSize * sizeof(int), stream)); + FLASHINFER_CUDA_CALL(cudaMemsetAsync( + sendRankLocalIndices, -1, maxTokenCountPerRank * maxSendRanksPerToken * sizeof(int), stream)); + FLASHINFER_CUDA_CALL(cudaMemsetAsync( + recvRankLocalIndices, -1, maxTokenCountPerRank * worldInfo.epSize * sizeof(int), stream)); + FLASHINFER_CUDA_CALL(cudaMemsetAsync(backwardRecvRankLocalIndices, -1, + maxTokenCountPerRank * maxSendRanksPerToken * sizeof(int), + stream)); + computeSendRecvRankCount(worldInfo, expertParallelInfo, maxTokenCountPerRank, + realRankTokenCountCumSum, gatheredTargetRankIds, sendRankCountCumSum, + recvRankCountCumSum, stream); + + inplaceSendRecvRankCumSum(worldInfo, sendRankCountCumSum, recvRankCountCumSum, stream); + computeSendRecvIndices(worldInfo, expertParallelInfo, maxTokenCountPerRank, + realRankTokenCountCumSum, gatheredTargetRankIds, sendRankCountCumSum, + recvRankCountCumSum, localGatherIndices, sendRankLocalIndices, + recvRankLocalIndices, backwardRecvRankLocalIndices, stream); + return cudaSuccess; +} + +template +__global__ void moeLocalGatherDevice(MoeEpWorldInfo worldInfo, + MoeExpertParallelInfo expertParallelInfo, + int maxTokenCountPerRank, int localMaxTokenCount, + int const* recvRankCountCumSum, int const* localGatherIndices, + int const* gatheredExpertIds, float const* gatheredScales, + int* localExpertIds, float* localScales) { + cg::thread_block_tile tile = + cg::tiled_partition(cg::this_thread_block()); + int laneInTile = tile.thread_rank(); + int tileId = threadIdx.x / kThreadsGroupSize; + int tileCountPerBlock = blockDim.x / kThreadsGroupSize; + + int epSize = worldInfo.epSize; + int rankTokenCount = recvRankCountCumSum[epSize - 1]; + bool needLoad = laneInTile < expertParallelInfo.topK; + + for (int index = tileId + blockIdx.x * tileCountPerBlock; index < localMaxTokenCount; + index += tileCountPerBlock * gridDim.x) { + int localTokenIndice = localGatherIndices[index]; + int expertId = needLoad && (index < rankTokenCount) + ? gatheredExpertIds[localTokenIndice * expertParallelInfo.topK + laneInTile] + : expertParallelInfo.expertCount; + float scale = needLoad && (index < rankTokenCount) + ? gatheredScales[localTokenIndice * expertParallelInfo.topK + laneInTile] + : 0.0f; + if (needLoad) { + localExpertIds[index * expertParallelInfo.topK + laneInTile] = expertId; + localScales[index * expertParallelInfo.topK + laneInTile] = scale; + } + } +} + +void moeLocalGather(MoeEpWorldInfo worldInfo, MoeExpertParallelInfo expertParallelInfo, + int maxTokenCountPerRank, int localMaxTokenCount, + int const* recvRankCountCumSum, int const* localGatherIndices, + int const* gatheredExpertIds, float const* gatheredScales, int* localExpertIds, + float* localScales, cudaStream_t stream) { + FLASHINFER_CHECK(expertParallelInfo.topK <= 32, + "Only topK less than or equal to 32 supported now."); + auto* kernelPtr = moeLocalGatherDevice<32>; + int paddedTopK = 32; + if (expertParallelInfo.topK <= 1) { + paddedTopK = 1; + kernelPtr = moeLocalGatherDevice<1>; + } else if (expertParallelInfo.topK <= 2) { + paddedTopK = 2; + kernelPtr = moeLocalGatherDevice<2>; + } else if (expertParallelInfo.topK <= 4) { + paddedTopK = 4; + kernelPtr = moeLocalGatherDevice<4>; + } else if (expertParallelInfo.topK <= 8) { + paddedTopK = 8; + kernelPtr = moeLocalGatherDevice<8>; + } else if (expertParallelInfo.topK <= 16) { + paddedTopK = 16; + kernelPtr = moeLocalGatherDevice<16>; + } + + int threadsPerBlock = 512; + int tokenPerBlock = threadsPerBlock / paddedTopK; + int blockCount = (localMaxTokenCount + tokenPerBlock - 1) / tokenPerBlock * 2; + + kernelPtr<<>>( + worldInfo, expertParallelInfo, maxTokenCountPerRank, localMaxTokenCount, recvRankCountCumSum, + localGatherIndices, gatheredExpertIds, gatheredScales, localExpertIds, localScales); +} + +int AllToAllChannelCommunicatorBase::maxSmCount = -1; +bool AllToAllChannelCommunicatorBase::maxSmCountUsed = false; + +void setMaxUsableSmCount(int smCount) { + AllToAllChannelCommunicatorBase::setMaxUsableSmCount(smCount); +} + +} // namespace trtllm_alltoall +} // namespace flashinfer diff --git a/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh b/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh new file mode 100644 index 0000000000000000000000000000000000000000..3dbed4b649b36c2a37bf39dd26d216fcd52ced8a --- /dev/null +++ b/include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh @@ -0,0 +1,675 @@ +/* + * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include + +#include + +#include "../exception.h" +#include "../logging.h" +namespace flashinfer { +namespace trtllm_mnnvl_allreduce { + +template +struct AllReduceParams { + int nranks; + int rank; + int buffer_M; + int num_tokens; + int token_dim; + void** buffer_ptrs_dev; + void* multicast_ptr; + void* buffer_flags; + bool wait_for_results; + bool launch_with_pdl; + + void* input; + void* output; + cudaStream_t stream; +}; + +template +struct RMSNormParams { + void* residual_output; + void* output; + void const* input; + void const* gamma; + double epsilon; + void* residual; + uint32_t* buffer_flags; + int batch; + int hidden_dim; + cudaStream_t stream; + bool launch_with_pdl; +}; + +__device__ bool isNegZero(float v) { return v == 0.f && signbit(v); } + +__device__ bool isNegZero(__nv_bfloat16 val) { return isNegZero(__bfloat162float(val)); } + +__device__ bool isNegZero(__nv_half val) { return isNegZero(__half2float(val)); } + +template +inline __device__ float toFloat(T val) { + return val; +} + +template <> +inline __device__ float toFloat<__nv_bfloat16>(__nv_bfloat16 val) { + return __bfloat162float(val); +} + +template <> +inline __device__ float toFloat<__nv_half>(__nv_half val) { + return __half2float(val); +} + +template +inline __device__ T fromFloat(float val) { + return val; +} + +template <> +inline __device__ __nv_bfloat16 fromFloat<__nv_bfloat16>(float val) { + return __float2bfloat16(val); +} + +template <> +inline __device__ __nv_half fromFloat<__nv_half>(float val) { + return __float2half(val); +} + +inline __device__ float2 loadfloat2(void const* ptr) { + float2 return_value; + asm volatile("ld.volatile.global.v2.f32 {%0, %1}, [%2];\n" + : "=f"(return_value.x), "=f"(return_value.y) + : "l"(ptr)); + return return_value; +} + +template +inline __device__ T divUp(T val, T divisor) { + return (val + divisor - 1) / divisor; +} + +__device__ struct __attribute__((aligned(32))) LamportFlags { + uint32_t buffer_size; + uint32_t input_offset; + uint32_t clear_offset; + uint32_t num_tokens_prev; + uint32_t* offset_access_ptr; + uint32_t* buffer_flags; + + __device__ explicit LamportFlags(uint32_t* buffer_flags) + : offset_access_ptr(&buffer_flags[4]), buffer_flags(buffer_flags) { + uint4 flag = reinterpret_cast(buffer_flags)[0]; + buffer_size = flag.z; + input_offset = flag.x * (buffer_size << 1U); + clear_offset = flag.y * (buffer_size << 1U); + num_tokens_prev = flag.w; + } + + __device__ void cta_arrive() { + __syncthreads(); + if (threadIdx.x == 0) { +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)) + asm volatile("red.async.release.global.gpu.add.u32 [%0], %1;" ::"l"(offset_access_ptr), "r"(1) + : "memory"); +#elif (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("red.global.gpu.add.u32 [%0], %1;" ::"l"(offset_access_ptr), "r"(1) : "memory"); +#else + atomicAdd(offset_access_ptr, 1); +#endif + } + } + + __device__ void wait_and_update(uint32_t num_tokens) { + if (threadIdx.x == 0 && blockIdx.x == gridDim.x - 1 && blockIdx.y == 0) { + while (*reinterpret_cast(offset_access_ptr) < gridDim.x * gridDim.y) { + } + uint4 flag = reinterpret_cast(buffer_flags)[0]; + buffer_flags[0] = (flag.x + 1) % 3; + buffer_flags[1] = (flag.y + 1) % 3; + buffer_flags[3] = num_tokens; + *(offset_access_ptr) = 0; + } + } +}; + +template +__global__ void twoshot_allreduce_kernel(T* output_ptr, T* shard_ptr, T** input_ptrs, T* mcast_ptr, + int num_tokens, int buffer_M, int token_dim, int rank, + uint32_t* buffer_flags, bool wait_for_results) { + int elt = blockIdx.y * blockDim.x + threadIdx.x; + + if (elt >= token_dim) return; + int token = blockIdx.x; + +#if (__CUDACC_VER_MAJOR__ >= 12 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + cudaGridDependencySynchronize(); +#endif + + LamportFlags flags(buffer_flags); + + // Capture the number of tokens in previous iteration so that we can properly clear the buffer + // The scatter stage will use the buffer in WORLD_SIZE granularity, thus we need to round up + uint32_t clr_toks_cta = + divUp(flags.num_tokens_prev > num_tokens ? flags.num_tokens_prev : num_tokens, + WORLD_SIZE) * + WORLD_SIZE; + clr_toks_cta = divUp(clr_toks_cta, gridDim.x); + + if (elt < token_dim) { + // Scatter token + int dest_rank = token % WORLD_SIZE; + int dest_token_offset = token / WORLD_SIZE; + T val = shard_ptr[token * token_dim + elt]; + if (isNegZero(val)) val = fromFloat(0.f); + input_ptrs[dest_rank][flags.input_offset + dest_token_offset * token_dim * WORLD_SIZE + + rank * token_dim + elt] = val; + + // Clear the buffer used by the previous call. Note the number of tokens to clear could be + // larger than the + // number of tokens in the current call. + for (int clr_tok = 0; clr_tok < clr_toks_cta; clr_tok++) { + uint32_t clr_token_idx = token + clr_tok * gridDim.x; + if (clr_token_idx < buffer_M) { + input_ptrs[rank][flags.clear_offset + clr_token_idx * token_dim + elt] = fromFloat(-0.f); + } + } + + // Reduce and broadcast + if ((token % WORLD_SIZE) == rank) { + int local_token = token / WORLD_SIZE; + float accum = 0.f; + + T values[WORLD_SIZE]; + + while (1) { + bool valid = true; + for (int r = 0; r < WORLD_SIZE; r++) { + T volatile* lamport_ptr = + (T volatile*)&input_ptrs[rank] + [flags.input_offset + local_token * token_dim * WORLD_SIZE + + r * token_dim + elt]; + values[r] = *lamport_ptr; + valid &= !isNegZero(values[r]); + } + if (valid) break; + } + for (int r = 0; r < WORLD_SIZE; r++) { + accum += toFloat(values[r]); + } + mcast_ptr[flags.input_offset + buffer_M * token_dim + token * token_dim + elt] = + fromFloat(accum); + } + } + +#if (__CUDACC_VER_MAJOR__ >= 12 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + cudaTriggerProgrammaticLaunchCompletion(); +#endif + + // Similarly clear broadcast buffer here + for (int clr_tok = 0; clr_tok < clr_toks_cta; clr_tok++) { + uint32_t clr_token_idx = token + clr_tok * gridDim.x; + if (clr_token_idx < buffer_M) { + input_ptrs[rank][flags.clear_offset + buffer_M * token_dim + clr_token_idx * token_dim + + elt] = fromFloat(-0.f); + } + } + + // Optionally wait for results if the next layer isn't doing the Lamport check + if (wait_for_results) { + // Update the atomic counter to indicate the block has read the offsets + flags.cta_arrive(); + // Only use a set of CTAs for lamport sync, reargange the grid + constexpr int ELTS_PER_LOAD = sizeof(float2) / sizeof(T); + // blockDim.x / ELTS_PER_LOAD should be at least the size of a warp (32) + if (threadIdx.x < (blockDim.x / ELTS_PER_LOAD)) { + uint64_t current_pos = + blockIdx.x * token_dim + blockIdx.y * blockDim.x + threadIdx.x * ELTS_PER_LOAD; + + void* lamport_ptr = + (void*)&input_ptrs[rank][flags.input_offset + buffer_M * token_dim + current_pos]; + // We have 2 assumptions here: + // 1. The write is atomic in 8B granularity -> Each buffer in the buffer group should be + // aligned to 8B + // 2. The num_token * token_dim is divisible by ELTS_PER_LOAD (4 for BF16 and 2 for FP32) + float2 val = loadfloat2(lamport_ptr); + while (isNegZero(*(T*)&val)) { + val = loadfloat2(lamport_ptr); + } + if (output_ptr) { + *((float2*)&output_ptr[current_pos]) = val; + } + } + + // Update the buffer flags + flags.wait_and_update(num_tokens); + } +} + +// Template-based dispatch functions following the same pattern as trtllm_allreduce.cuh +template +cudaError_t twoshot_allreduce_dispatch(AllReduceParams& params) { + int const num_threads = 128; + int const num_blocks = (params.token_dim + num_threads - 1) / num_threads; + + dim3 grid(params.num_tokens, num_blocks); + + cudaLaunchConfig_t config; + cudaLaunchAttribute attrs[1]; + config.dynamicSmemBytes = 0; + config.stream = params.stream; + config.gridDim = grid; + config.blockDim = num_threads; + config.attrs = attrs; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = params.launch_with_pdl ? 1 : 0; + config.numAttrs = 1; + + cudaLaunchKernelEx(&config, &twoshot_allreduce_kernel, + reinterpret_cast(params.output), reinterpret_cast(params.input), + reinterpret_cast(params.buffer_ptrs_dev), + reinterpret_cast(params.multicast_ptr), params.num_tokens, params.buffer_M, + params.token_dim, params.rank, + reinterpret_cast(params.buffer_flags), params.wait_for_results); + + return cudaSuccess; +} + +template +cudaError_t twoshot_allreduce_dispatch_world_size(AllReduceParams& params) { + FLASHINFER_LOG_DEBUG("twoshot_allreduce_dispatch_world_size"); + switch (params.nranks) { + case 2: + return twoshot_allreduce_dispatch(params); + case 4: + return twoshot_allreduce_dispatch(params); + case 8: + return twoshot_allreduce_dispatch(params); + case 16: + return twoshot_allreduce_dispatch(params); + case 32: + return twoshot_allreduce_dispatch(params); + case 64: + return twoshot_allreduce_dispatch(params); + default: + FLASHINFER_ERROR("MNNVL AllReduce: unsupported world_size " + std::to_string(params.nranks) + + ". Supported sizes: {2, 4, 8, 16, 32, 64}"); + return cudaErrorInvalidValue; + } +} + +template +__device__ void copy_f4(T_IN* dst, T_IN const* src) { + float4* dst4 = (float4*)dst; + float4 const* src4 = (float4 const*)src; + __pipeline_memcpy_async(dst4, src4, sizeof(float4)); +} + +template +__device__ void copy_f4_ldg(T_IN* dst, T_IN const* src) { + float4* dst4 = (float4*)dst; + float4 const* src4 = (float4*)src; + *dst4 = *src4; +} + +__device__ float4 loadfloat4(void const* ptr) { + // Check alignment - ptr should be 16-byte aligned for safe float4 load + if (reinterpret_cast(ptr) % 16 != 0) { + // Fall back to scalar loads if not aligned + float4 return_value; + float const* float_ptr = reinterpret_cast(ptr); + return_value.x = float_ptr[0]; + return_value.y = float_ptr[1]; + return_value.z = float_ptr[2]; + return_value.w = float_ptr[3]; + return return_value; + } + + float4 return_value; + + asm volatile("ld.volatile.global.v4.f32 {%0, %1, %2, %3}, [%4];\n" + : "=f"(return_value.x), "=f"(return_value.y), "=f"(return_value.z), + "=f"(return_value.w) + : "l"(ptr)); + + return return_value; +} + +// Safer version that checks bounds before loading +template +__device__ float4 loadfloat4_safe(T const* ptr, int remaining_elements) { + float return_value[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + + if (remaining_elements <= 0) { + return *(float4*)return_value; + } + + // Check alignment - ptr should be 16-byte aligned for safe float4 load + bool is_aligned = (reinterpret_cast(ptr) % 16 == 0); + + if (is_aligned && remaining_elements >= 4) { + // Safe to do vectorized load + asm volatile("ld.volatile.global.v4.f32 {%0, %1, %2, %3}, [%4];\n" + : "=f"(return_value[0]), "=f"(return_value[1]), "=f"(return_value[2]), + "=f"(return_value[3]) + : "l"(ptr)); + } else { + // Fall back to scalar loads with bounds checking + float const* float_ptr = reinterpret_cast(ptr); + for (int i = 0; i < 4 && i < remaining_elements; i++) { + return_value[i] = toFloat(float_ptr[i]); + } + } + + return *(float4*)return_value; +} + +template +inline __device__ T add(T a, T b) { + return a + b; +} + +#define FINAL_MASK 0xffffffff + +template +__inline__ __device__ T warpReduceSum(T val) { +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) + val = add(val, __shfl_xor_sync(FINAL_MASK, val, mask, + 32)); //__shfl_sync bf16 return float when sm < 80 + return val; +} + +inline __device__ float block_reduce_sum(float val) { + __shared__ float smem[32]; + int lane_id = threadIdx.x % 32, warp_id = threadIdx.x / 32, warp_num = blockDim.x / 32; + val = warpReduceSum(val); + if (lane_id == 0) { + smem[warp_id] = val; + } + __syncthreads(); + val = lane_id < warp_num ? smem[lane_id] : 0.f; + val = warpReduceSum(val); + return val; +} + +template +__global__ void __launch_bounds__(128, 1) + RMSNorm(T_IN* input_plus_residual, T_OUT* output_norm, T_IN const* buffer_input, + T_IN const* gamma, float epsilon, T_IN const* residual, int batch_size, + uint32_t* buffer_flags) { +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + + static bool const LAMPORT = true; + + extern __shared__ uint8_t smem[]; + + int sample = blockIdx.y; + + static int const CGA_THREADS = NUM_THREADS * 1; + + static int const ITERS = DIM / CGA_THREADS; + float r_input[ITERS]; + float r_gamma[ITERS]; + + T_IN* sh_input = (T_IN*)&smem[0]; + T_IN* sh_residual = (T_IN*)&smem[NUM_INPUTS * NUM_THREADS * ITERS * sizeof(T_IN)]; + T_IN* sh_gamma = (T_IN*)&smem[(NUM_INPUTS + 1) * NUM_THREADS * ITERS * sizeof(T_IN)]; + + static int const ELTS_PER_THREAD = sizeof(float4) / sizeof(T_IN); + + int offsets[NUM_INPUTS][DIM / (1 * ELTS_PER_THREAD * NUM_THREADS)]; + + LamportFlags flags(buffer_flags); + T_IN const* input = &buffer_input[flags.input_offset + flags.buffer_size]; + +#if (__CUDACC_VER_MAJOR__ >= 12 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + cudaTriggerProgrammaticLaunchCompletion(); +#endif + + for (int i = 0; i < NUM_INPUTS; i++) { + for (int j = 0; j < DIM / (1 * ELTS_PER_THREAD * NUM_THREADS); j++) { + int k = j * NUM_THREADS + threadIdx.x; + offsets[i][j] = + i * batch_size * DIM + sample * DIM + blockIdx.x * DIM / 1 + k * ELTS_PER_THREAD; + } + } + +#pragma unroll + for (int j = 0; j < DIM / (1 * ELTS_PER_THREAD * NUM_THREADS); j++) { + int i = j * NUM_THREADS + threadIdx.x; + copy_f4(&sh_residual[i * ELTS_PER_THREAD], + &residual[sample * DIM + blockIdx.x * DIM + i * ELTS_PER_THREAD]); + } + + __pipeline_commit(); + +#pragma unroll + for (int j = 0; j < DIM / (ELTS_PER_THREAD * NUM_THREADS); j++) { + int i = j * NUM_THREADS + threadIdx.x; + copy_f4(&sh_gamma[i * ELTS_PER_THREAD], &gamma[blockIdx.x * DIM + i * ELTS_PER_THREAD]); + } + + __pipeline_commit(); + flags.cta_arrive(); + + // Load all inputs + bool valid = false; + +#if (__CUDACC_VER_MAJOR__ >= 12 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + if (!LAMPORT) cudaGridDependencySynchronize(); +#endif + + while (!valid) { + valid = true; +#pragma unroll + for (int i = 0; i < NUM_INPUTS; i++) { + for (int j = 0; j < DIM / (ELTS_PER_THREAD * NUM_THREADS); j++) { + int k = j * NUM_THREADS + threadIdx.x; + + float4* dst4 = (float4*)&sh_input[i * NUM_THREADS * ITERS + k * ELTS_PER_THREAD]; + + // Calculate the absolute element offset from the start of buffer_input + int element_offset = offsets[i][j]; + + // The input pointer is already offset to: &buffer_input[buffer_offset + buffer_size] + // So the actual pointer we're accessing is: input + element_offset + // Which equals: &buffer_input[buffer_offset + buffer_size + element_offset] + + float4* src4 = (float4*)&input[element_offset]; + + float4 value; + // Check if we have enough elements remaining for a safe float4 load + if (element_offset >= 0 && element_offset + ELTS_PER_THREAD <= flags.buffer_size) { + value = loadfloat4(src4); + } else { + // Use safe load for boundary cases or out-of-bounds + int remaining_elements = flags.buffer_size - element_offset; + if (remaining_elements <= 0) { + // Completely out of bounds, return zeros + float4 return_value = {0.0f, 0.0f, 0.0f, 0.0f}; + value = return_value; + } else { + value = loadfloat4_safe(reinterpret_cast(src4), remaining_elements); + } + } + + if (LAMPORT) { + // Assume that the 16B were written atomically, so we only need to check one value + T_IN lowest_val = *(T_IN*)&value; + valid &= !isNegZero(lowest_val); + } + *dst4 = value; + } + } + } + + __syncthreads(); + + // Perform the initial input reduction + if (NUM_INPUTS > 0) { + T_IN accum[ELTS_PER_THREAD]; + float4* accum4 = (float4*)&accum; + + for (int j = 0; j < DIM / (ELTS_PER_THREAD * NUM_THREADS); j++) { + int k = j * NUM_THREADS + threadIdx.x; + + *accum4 = *(float4*)&sh_input[k * ELTS_PER_THREAD]; + + for (int i = 1; i < NUM_INPUTS; i++) { + float4 data = *(float4*)&sh_input[i * NUM_THREADS * ITERS + k * ELTS_PER_THREAD]; + T_IN* p_d = (T_IN*)&data; + for (int x = 0; x < ELTS_PER_THREAD; x++) { + accum[x] += p_d[x]; + } + } + + // Write back to input 0's staging location. No sync needed since all data localized to + // thread. + *(float4*)&sh_input[k * ELTS_PER_THREAD] = *accum4; + } + } + + // Wait for residual + __pipeline_wait_prior(1); + __syncthreads(); + + float thread_sum = 0.f; + +#pragma unroll + for (int io = 0; io < ITERS / ELTS_PER_THREAD; io++) { + float4 inp4 = + *(float4*)&sh_input[io * NUM_THREADS * ELTS_PER_THREAD + threadIdx.x * ELTS_PER_THREAD]; + float4 res4 = + *(float4*)&sh_residual[io * NUM_THREADS * ELTS_PER_THREAD + threadIdx.x * ELTS_PER_THREAD]; + + T_IN* r_inp = (T_IN*)&inp4; + T_IN* r_res = (T_IN*)&res4; + + float4 out4; + + T_IN* r_out = (T_IN*)&out4; + + for (int ii = 0; ii < ELTS_PER_THREAD; ii++) { + int i = io * ELTS_PER_THREAD + ii; + + T_IN inp_plus_resid = r_inp[ii] + r_res[ii]; + r_out[ii] = inp_plus_resid; + r_input[i] = toFloat(inp_plus_resid); + + // Accumulate the squares for RMSNorm + thread_sum += toFloat(inp_plus_resid * inp_plus_resid); + } + + *(float4*)&input_plus_residual[sample * DIM + blockIdx.x * DIM + + io * NUM_THREADS * ELTS_PER_THREAD + + threadIdx.x * ELTS_PER_THREAD] = out4; + } + + // Wait for Gamma. There will be a global synchronization as part of the reduction + __pipeline_wait_prior(0); + + float cluster_sum = block_reduce_sum(thread_sum); + + float rcp_rms = rsqrtf(cluster_sum / DIM + epsilon); + +#pragma unroll + for (int io = 0; io < ITERS / ELTS_PER_THREAD; io++) { + float4 gamma4 = + *(float4*)&sh_gamma[io * NUM_THREADS * ELTS_PER_THREAD + threadIdx.x * ELTS_PER_THREAD]; + T_IN* r_g4 = (T_IN*)&gamma4; + + float4 out4; + // FIXME: this only works if T_OUT == T_IN + T_OUT* r_out = (T_OUT*)&out4; + + for (int ii = 0; ii < ELTS_PER_THREAD; ii++) { + int i = io * ELTS_PER_THREAD + ii; + r_gamma[i] = toFloat(r_g4[ii]); + r_out[ii] = fromFloat(r_gamma[i] * r_input[i] * rcp_rms); + } + + *(float4*)&output_norm[sample * DIM + blockIdx.x * DIM + io * NUM_THREADS * ELTS_PER_THREAD + + threadIdx.x * ELTS_PER_THREAD] = out4; + } + // Update the buffer pointers + flags.wait_and_update(batch_size); +#endif +} + +template +cudaError_t twoshot_rmsnorm_dispatch(RMSNormParams& params) { + static constexpr int NUM_THREADS = 128; + static constexpr int CGA_THREADS = NUM_THREADS; + constexpr int iters = H_DIM / CGA_THREADS; + + dim3 grid(1, params.batch, 1); + + cudaLaunchConfig_t config; + cudaLaunchAttribute attrs[1]; + config.stream = params.stream; + config.gridDim = grid; + config.blockDim = NUM_THREADS; + config.attrs = attrs; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = params.launch_with_pdl ? 1 : 0; + config.numAttrs = 1; + + size_t shmem_size = 3 * NUM_THREADS * iters * sizeof(T); + config.dynamicSmemBytes = shmem_size; + + cudaFuncSetAttribute(&RMSNorm, + cudaFuncAttributeMaxDynamicSharedMemorySize, shmem_size); + + cudaLaunchKernelEx( + &config, &RMSNorm, reinterpret_cast(params.residual_output), + reinterpret_cast(params.output), reinterpret_cast(params.input), + reinterpret_cast(params.gamma), static_cast(params.epsilon), + reinterpret_cast(params.residual), params.batch, params.buffer_flags); + + return cudaSuccess; +} + +template +cudaError_t twoshot_rmsnorm_dispatch_hidden_dim(RMSNormParams& params) { + FLASHINFER_LOG_DEBUG("twoshot_rmsnorm_dispatch_hidden_dim"); + switch (params.hidden_dim) { + case 2048: + return twoshot_rmsnorm_dispatch(params); + case 4096: + return twoshot_rmsnorm_dispatch(params); + case 5120: + return twoshot_rmsnorm_dispatch(params); // Llama-4 + case 7168: + return twoshot_rmsnorm_dispatch(params); // DeepSeek + case 8192: + return twoshot_rmsnorm_dispatch(params); + default: + FLASHINFER_ERROR("MNNVL TwoShot RMSNorm: unsupported hidden_dim " + + std::to_string(params.hidden_dim) + + ". Supported sizes: {2048, 4096, 5120, 7168, 8192}"); + return cudaErrorInvalidValue; + } +} + +} // namespace trtllm_mnnvl_allreduce +} // namespace flashinfer diff --git a/include/flashinfer/comm/trtllm_moe_allreduce_fusion.cuh b/include/flashinfer/comm/trtllm_moe_allreduce_fusion.cuh new file mode 100644 index 0000000000000000000000000000000000000000..6037c1979897c61f4eba08546bf3da28ad6fb965 --- /dev/null +++ b/include/flashinfer/comm/trtllm_moe_allreduce_fusion.cuh @@ -0,0 +1,1496 @@ +#include +#include +#include +#include + +#include +#include +#include + +#include "../exception.h" +#include "../logging.h" +#include "../utils.cuh" +#include "../vec_dtypes.cuh" + +namespace flashinfer { + +namespace trtllm_moe_allreduce_fusion { + +namespace details { + +static constexpr int CVT_FP4_ELTS_PER_THREAD = 8; +static constexpr int CVT_FP4_SF_VEC_SIZE = 16; +static constexpr int kBytesPerAccess = 16; +static constexpr int kOneShotMaxToken = 128; +static constexpr int kBarrierFlagCount = 256; + +} // namespace details + +namespace maths { +// // ============================== Cast ============================== +template +__device__ inline T_OUT cuda_cast(T_IN val) { + return val; +} + +template <> +__device__ inline float2 cuda_cast(int2 val) { + return make_float2(val.x, val.y); +} + +template <> +__device__ inline float2 cuda_cast(float val) { + return make_float2(val, val); +} + +template <> +__device__ inline float2 cuda_cast(half2 val) { + return __half22float2(val); +} + +template <> +__device__ inline half2 cuda_cast(float2 val) { + return __float22half2_rn(val); +} + +template <> +__device__ inline half2 cuda_cast(float val) { + return __float2half2_rn(val); +} + +template <> +__device__ inline half2 cuda_cast(half val) { + return __half2half2(val); +} + +template <> +__device__ inline int8_t cuda_cast(half val) { + union { + int8_t int8[2]; + int16_t int16; + }; + + union { + half fp16; + int16_t int16_in; + }; + + fp16 = val; + asm volatile("cvt.rni.sat.s8.f16 %0, %1;" : "=h"(int16) : "h"(int16_in)); + return int8[0]; +} + +template <> +__device__ inline int16_t cuda_cast(half2 val) { + union { + int8_t int8[2]; + int16_t int16; + }; + + int8[0] = cuda_cast(val.x); + int8[1] = cuda_cast(val.y); + return int16; +} + +template <> +__device__ inline int8_t cuda_cast(float val) { + union { + int8_t int8[2]; + int16_t int16; + }; + + asm volatile("cvt.rni.sat.s8.f32 %0, %1;" : "=h"(int16) : "f"(val)); + return int8[0]; +} + +template <> +__device__ inline int16_t cuda_cast(float2 val) { + union { + int8_t int8[2]; + int16_t int16; + }; + + int8[0] = cuda_cast(val.x); + int8[1] = cuda_cast(val.y); + return int16; +} + +template <> +__device__ inline half2 cuda_cast(int16_t val) { + union { + int8_t int8[2]; + int16_t int16; + }; + + int16 = val; + return make_half2(int8[0], int8[1]); +} + +template <> +__device__ inline float2 cuda_cast(int16_t val) { + union { + int8_t int8[2]; + int16_t int16; + }; + + int16 = val; + return make_float2(int8[0], int8[1]); +} + +template <> +__device__ inline __nv_bfloat16 cuda_cast(int32_t val) { + return static_cast(val); +} + +template <> +__device__ inline __nv_bfloat16 cuda_cast(int8_t val) { + return static_cast(val); +} + +template <> +__device__ inline int8_t cuda_cast(__nv_bfloat16 val) { + return static_cast(val); +} + +template <> +__device__ inline float cuda_cast(__nv_bfloat16 val) { + return __bfloat162float(val); +} + +inline __device__ float2 bf1622float2(const __nv_bfloat162 val) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + float2 f_val; + f_val.x = __low2float(val); + f_val.y = __high2float(val); + return f_val; +#else + return __bfloat1622float2(val); +#endif +} + +template <> +__device__ inline float2 cuda_cast(__nv_bfloat162 val) { + return bf1622float2(val); +} + +template <> +__device__ inline half cuda_cast(__nv_bfloat16 val) { + return __float2half(__bfloat162float(val)); +} + +inline __device__ int16_t bf1622int16(__nv_bfloat162 val) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + float2 f_val; + f_val.x = max(min(__low2float(val), 127.f), -128.f); + f_val.y = max(min(__high2float(val), 127.f), -128.f); + + union { + int8_t int8[2]; + int16_t int16; + }; + + int8[0] = static_cast(static_cast(f_val.x)); + int8[1] = static_cast(static_cast(f_val.y)); + return int16; +#else + val = __hmin2(val, make_bfloat162(127., 127.)); + val = __hmax2(val, make_bfloat162(-128., -128.)); + + union { + int8_t int8[2]; + int16_t int16; + }; + + int8[0] = static_cast(static_cast(val.x)); + int8[1] = static_cast(static_cast(val.y)); + return int16; +#endif +} + +template <> +__device__ inline int16_t cuda_cast(__nv_bfloat162 val) { + return bf1622int16(val); +} + +template <> +__device__ inline __nv_bfloat16 cuda_cast<__nv_bfloat16, float>(float val) { + return __float2bfloat16(val); +} + +template <> +__device__ inline __nv_bfloat16 cuda_cast<__nv_bfloat16, half>(half val) { + return __float2bfloat16(__half2float(val)); +} + +inline __device__ __nv_bfloat162 bf162bf162(const __nv_bfloat16 val) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + __nv_bfloat162 val2; + val2.x = val; + val2.y = val; + return val2; +#else + return __bfloat162bfloat162(val); +#endif +} + +template <> +__device__ inline __nv_bfloat162 cuda_cast<__nv_bfloat162, __nv_bfloat16>(__nv_bfloat16 val) { + return bf162bf162(val); +} + +template <> +__device__ inline __nv_bfloat162 cuda_cast<__nv_bfloat162, float>(float val) { + return __float2bfloat162_rn(val); +} + +inline __device__ __nv_bfloat162 float22bf162(const float2 val) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + return __floats2bfloat162_rn(val.x, val.y); +#else + return __float22bfloat162_rn(val); +#endif +} + +template <> +__device__ inline __nv_bfloat162 cuda_cast<__nv_bfloat162, float2>(float2 val) { + return float22bf162(val); +} + +template <> +__device__ inline __nv_bfloat162 cuda_cast<__nv_bfloat162, int16_t>(int16_t val) { + union { + int8_t int8[2]; + int16_t int16; + }; + + int16 = val; + __nv_bfloat162 res; + res.x = cuda_cast<__nv_bfloat16>(int8[0]); + res.y = cuda_cast<__nv_bfloat16>(int8[1]); + return res; +} + +template <> +__device__ inline __nv_bfloat162 cuda_cast<__nv_bfloat162, half2>(half2 val) { + return float22bf162(__half22float2(val)); +} + +// // ============================== Abs ============================== +template +__device__ inline T cuda_abs(T val) { + assert(false); + return {}; +} + +template <> +__device__ inline float cuda_abs(float val) { + return fabs(val); +} + +template <> +__device__ inline float2 cuda_abs(float2 val) { + return make_float2(fabs(val.x), fabs(val.y)); +} + +template <> +__device__ inline half cuda_abs(half val) { + return __habs(val); +} + +template <> +__device__ inline half2 cuda_abs(half2 val) { + return __habs2(val); +} + +#if __CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__) +template <> +__device__ inline __nv_bfloat16 cuda_abs(__nv_bfloat16 val) { + return __habs(val); +} + +template <> +__device__ inline __nv_bfloat162 cuda_abs(__nv_bfloat162 val) { + return __habs2(val); +} +#endif + +// // ============================== Max ============================== +template +__device__ inline To cuda_max(Ti val) { + return cuda_cast(val); +}; + +template <> +__device__ inline float cuda_max(float2 val) { + return fmaxf(val.x, val.y); +} + +template <> +__device__ inline half cuda_max(half2 val) { + return __hmax(val.x, val.y); +} + +template <> +__device__ inline __nv_bfloat16 cuda_max(__nv_bfloat162 val) { +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)) + return __hmax(val.x, val.y); +#else + assert(0); + asm volatile("brkpt;\n" ::); + return __nv_bfloat16(0); +#endif +} + +// Binary maximum: compute the max of two values. +template +__device__ inline T cuda_max(T val1, T val2) { + return (val1 > val2) ? val1 : val2; +} + +template <> +__device__ inline float2 cuda_max(float2 val1, float2 val2) { + float2 out; + out.x = fmaxf(val1.x, val2.x); + out.y = fmaxf(val1.y, val2.y); + return out; +} + +template <> +__device__ inline half2 cuda_max(half2 val1, half2 val2) { + return __hmax2(val1, val2); +} + +template <> +__device__ inline __nv_bfloat162 cuda_max(__nv_bfloat162 val1, __nv_bfloat162 val2) { + return __hmax2(val1, val2); +} + +// // ============================== Reciprocal ============================== +// Fast reciprocal. +inline __device__ float reciprocal_approximate_ftz(float a) { + float b; + asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a)); + return b; +} +} // namespace maths + +enum class FP4QuantizationSFLayout { + // Block scale factors are stored in swizzled layout for cutlass FP4 kernel. Scale factor + // blocks are organized in 512-byte blocks in global memory, with each block having 128x4 FP8 + // values. The SF matrix dimensions are therefore padded - rows to the nearest multiple of 128 and + // columns to the nearest multiple of 4. + // + // The scale factor block rows map to data block rows in an interleaved pattern: + // For a scale factor row 'i', it maps to data block row: (i % 4) * 32 + (i / 4) + // Column 'j' in the scale factor block corresponds to scaling the j-th block in the data tensor. + // + // Please refer to https://nvbugs/4165523 for more details about the swizzled layout. + SWIZZLED, + // Block scale factors are stored in linear layout (row-major). This is used in some trtllm-gen + // kernels standard. + LINEAR +}; + +namespace utils { +#define FINAL_MASK 0xffffffff + +template +__inline__ __device__ T warpReduceSumV2(T* val) { +#pragma unroll + for (int i = 0; i < NUM; i++) { +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) + val[i] += __shfl_xor_sync(FINAL_MASK, val[i], mask, 32); + } + return (T)(0.0f); +} + +template +__inline__ __device__ T blockReduceSumV2(T* val) { + static __shared__ T shared[NUM][33]; + int lane = threadIdx.x & 0x1f; + int wid = threadIdx.x >> 5; + + warpReduceSumV2(val); + + if (lane == 0) { +#pragma unroll + for (int i = 0; i < NUM; i++) { + shared[i][wid] = val[i]; + } + } + + __syncthreads(); + + bool is_mask = threadIdx.x < (blockDim.x / 32.f); +#pragma unroll + for (int i = 0; i < NUM; i++) { + val[i] = is_mask ? shared[i][lane] : (T)(0.0f); + } + warpReduceSumV2(val); + return (T)0.0f; +} + +inline __device__ int64_t get_sf_out_offset_128x4(std::optional batchIdx, int mIdx, int kIdx, + std::optional numRows, int numCols) { + // SF layout [numMTiles, numKTiles, 32 (mTile), 4 (mTile), 4(kTile)] + // --> index [mTileIdx, kTileIdx, outerMIdx, innerMIdx, innerKIdx] + + // batched tensor + // SF layout [numBTiles, numMTiles, numKTiles, 32 (mTile), 4 (mTile), 4(kTile)] + // --> index [bTileIdx, mTileIdx, kTileIdx, outerMIdx, innerMIdx, innerKIdx] + + int32_t innerKIdx = (kIdx % 4); + int64_t innerKStride = 1; + + int32_t innerMIdx = (mIdx % (32 * 4)) / 32; + int64_t innerMStride = 4 * innerKStride; // 4 + + // M tile layout [32, 4] is column-major. + int32_t outerMIdx = (mIdx % 32); + int64_t outerMStride = 4 * innerMStride; // 16 + + int32_t kTileIdx = (kIdx / 4); + int64_t kTileStride = 32 * outerMStride; // 512 + + // SF vector size 16. We round the "numCols" up to a multiple of 64. + int factor = details::CVT_FP4_SF_VEC_SIZE * 4; + int32_t numKTiles = (numCols + factor - 1) / factor; + int32_t mTileIdx = mIdx / (32 * 4); + int64_t mTileStride = numKTiles * kTileStride; + + // Each SF block has 128 rows so pad rows to the multiple of 128. + int32_t numMTiles = (numRows.value_or(0) + 128 - 1) / 128; + int64_t bTileStride = numMTiles * mTileStride; + + // Compute the global offset. + int64_t SFOffset = batchIdx.value_or(0) * bTileStride + mTileIdx * mTileStride + + kTileIdx * kTileStride + outerMIdx * outerMStride + innerMIdx * innerMStride + + innerKIdx * innerKStride; + + return SFOffset; +} + +template +__device__ uint8_t* cvt_quant_to_fp4_get_sf_out_offset(std::optional batchIdx, int rowIdx, + int colIdx, std::optional numRows, + int numCols, SFType* SFout, + FP4QuantizationSFLayout layout) { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + static_assert(CVT_FP4_NUM_THREADS_PER_SF == 1 || CVT_FP4_NUM_THREADS_PER_SF == 2); + + // One pair of threads write one SF to global memory. + // TODO: stage through smem for packed STG.32 + // is it better than STG.8 from 4 threads ? + if (threadIdx.x % CVT_FP4_NUM_THREADS_PER_SF == 0) { + if (layout == FP4QuantizationSFLayout::SWIZZLED) { + // SF vector index (16 elements share one SF in the K dimension). + // numRows and numCols are unpadded. + int32_t kIdx = colIdx / CVT_FP4_NUM_THREADS_PER_SF; + int32_t mIdx = rowIdx; + + auto SFOffset = get_sf_out_offset_128x4(batchIdx, mIdx, kIdx, numRows, numCols); + return reinterpret_cast(SFout) + SFOffset; + } else if (layout == FP4QuantizationSFLayout::LINEAR) { + // Linear row-major layout, no padding required. + int32_t KTileIdx = colIdx / CVT_FP4_NUM_THREADS_PER_SF; + + int32_t numKTiles = numCols / details::CVT_FP4_SF_VEC_SIZE; + int64_t mTileStride = numKTiles; + + int64_t BTileStride = numRows.value_or(0) * mTileStride; + + int64_t SFOffset = batchIdx.value_or(0) * BTileStride + rowIdx * mTileStride + KTileIdx; + return reinterpret_cast(SFout) + SFOffset; + } else { + return nullptr; + } + } +#endif + return nullptr; +} + +__forceinline__ __device__ uint32_t pack_bytes(uint8_t c0, uint8_t c1, uint8_t c2, uint8_t c3) { + uint32_t val0 = c0; + uint32_t val1 = c1; + uint32_t val2 = c2; + uint32_t val3 = c3; + + return (val3 << 24) | (val2 << 16) | (val1 << 8) | val0; +} + +// Convert 8 float32 values into 8 e2m1 values (represented as one uint32_t). +// NOTE:bypass sm_100 requirement by __nv_cvt_float2_to_fp4x2 +inline __device__ uint32_t fp32_vec_to_e2m1(float (&array)[8]) { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + uint32_t val; + asm volatile( + "{\n" + ".reg .b8 byte0;\n" + ".reg .b8 byte1;\n" + ".reg .b8 byte2;\n" + ".reg .b8 byte3;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte0, %2, %1;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte1, %4, %3;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte2, %6, %5;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte3, %8, %7;\n" + "mov.b32 %0, {byte0, byte1, byte2, byte3};\n" + "}" + : "=r"(val) + : "f"(array[0]), "f"(array[1]), "f"(array[2]), "f"(array[3]), "f"(array[4]), "f"(array[5]), + "f"(array[6]), "f"(array[7])); + return val; +#else + uint32_t val; + __nv_fp4x2_storage_t vals[4]; +#pragma unroll + for (int i = 0; i < 4; i++) { + vals[i] = __nv_cvt_float2_to_fp4x2(*(((float2*)array) + i), __NV_E2M1, cudaRoundNearest); + } + val = pack_bytes(vals[0], vals[1], vals[2], vals[3]); + return val; +#endif +} + +// Convert 4 float2 values into 8 e2m1 values (represented as one uint32_t). +inline __device__ uint32_t fp32_vec_to_e2m1(float2 (&array)[4]) { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + uint32_t val; + asm volatile( + "{\n" + ".reg .b8 byte0;\n" + ".reg .b8 byte1;\n" + ".reg .b8 byte2;\n" + ".reg .b8 byte3;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte0, %2, %1;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte1, %4, %3;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte2, %6, %5;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte3, %8, %7;\n" + "mov.b32 %0, {byte0, byte1, byte2, byte3};\n" + "}" + : "=r"(val) + : "f"(array[0].x), "f"(array[0].y), "f"(array[1].x), "f"(array[1].y), "f"(array[2].x), + "f"(array[2].y), "f"(array[3].x), "f"(array[3].y)); + return val; +#else + uint32_t val; + __nv_fp4x2_storage_t vals[4]; +#pragma unroll + for (int i = 0; i < 4; i++) { + vals[i] = __nv_cvt_float2_to_fp4x2(array[i], __NV_E2M1, cudaRoundNearest); + } + val = pack_bytes(vals[0], vals[1], vals[2], vals[3]); + return val; +#endif +} + +// Quantizes the provided PackedVec into the uint32_t output +template +__device__ uint32_t cvt_warp_fp16_to_fp4(vec_t& vec, float SFScaleVal, + uint8_t* SFout) { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + // Get absolute maximum values among the local 8 values. + auto localMax = maths::cuda_abs(get_vec2_element(vec, 0)); + +#pragma unroll + for (int i = 1; i < details::CVT_FP4_ELTS_PER_THREAD / 2; i++) { + localMax = maths::cuda_max(localMax, maths::cuda_abs(get_vec2_element(vec, i))); + } + + // Get the absolute maximum among all 16 values (two threads). + localMax = maths::cuda_max(__shfl_xor_sync(uint32_t(-1), localMax, 1), localMax); + // Get the final absolute maximum values. + float vecMax = float(maths::cuda_max(localMax.x, localMax.y)); + + // Get the SF (max value of the vector / max value of e2m1). + // maximum value of e2m1 = 6.0. + // TODO: use half as compute data type. + float SFValue = SFScaleVal * (vecMax * maths::reciprocal_approximate_ftz(6.0f)); + // 8 bits representation of the SF. + uint8_t fp8SFVal; + // Write the SF to global memory (STG.8). + if constexpr (UE8M0_SF) { +#if (__CUDACC_VER_MAJOR__ * 10000 + __CUDACC_VER_MINOR__ * 100 >= 120800) + __nv_fp8_e8m0 tmp; + tmp.__x = __nv_cvt_float_to_e8m0(SFValue, __NV_SATFINITE, cudaRoundPosInf); + SFValue = static_cast(tmp); + fp8SFVal = tmp.__x; +#else +#error "FP8 E8M0 support requires CUDA 12.8 or newer." +#endif + } else { + // Here SFValue is always positive, so E4M3 is the same as UE4M3. + __nv_fp8_e4m3 tmp = __nv_fp8_e4m3(SFValue); + fp8SFVal = tmp.__x; + SFValue = static_cast(tmp); + } + // Get the output scale. + // Recipe: final_scale = reciprocal(fp32(fp8(SFValue * SFScaleVal))) * reciprocal(SFScaleVal)) + float outputScale = SFValue != 0 ? maths::reciprocal_approximate_ftz( + SFValue * maths::reciprocal_approximate_ftz(SFScaleVal)) + : 0.0f; + + if (SFout) { + // Write the SF to global memory (STG.8). + *SFout = fp8SFVal; + } + + // Convert the input to float. + float2 fp2Vals[details::CVT_FP4_ELTS_PER_THREAD / 2]; + +#pragma unroll + for (int i = 0; i < details::CVT_FP4_ELTS_PER_THREAD / 2; i++) { + if constexpr (std::is_same_v) { + fp2Vals[i] = __half22float2(get_vec2_element(vec, i)); + } else { + fp2Vals[i] = __bfloat1622float2(get_vec2_element(vec, i)); + } + fp2Vals[i].x *= outputScale; + fp2Vals[i].y *= outputScale; + } + + // Convert to e2m1 values. + uint32_t e2m1Vec = fp32_vec_to_e2m1(fp2Vals); + + // Write the e2m1 values to global memory. + return e2m1Vec; +#else + return 0; +#endif +} +} // namespace utils + +template +struct AllReduceFusionParams { + int nranks; + int rank; + // size = token_num * hidden_dim + int size; + int hidden_dim; + void** workspace; + void* allreduce_in; + void* allreduce_out; + void* residual_in; + void* residual_out; + void* norm_out; + void* quant_out; + void* scale_out; + void* rms_gamma; + float rms_eps; + // todo(review): why float* scale_factor in trt-llm? + float scale_factor; + FP4QuantizationSFLayout layout = FP4QuantizationSFLayout::SWIZZLED; + cudaStream_t stream; + + // moe-allreduce output (non-fused) + // might be used in MoeReductionAllReduceFusionParams + void* moe_allreduce_out = nullptr; +}; + +template +struct MoeReductionAllReduceFusionParams : public AllReduceFusionParams { + // * moe reduction specific params + // Refer to kernel implementation on layout of those params + // number of active experts on current device + + int moe_reduction_device_num_experts = 0; + // per token per expert fp32 scale + float* moe_reduction_scale_input = nullptr; + // per token per expert input + void* moe_reduction_active_experts_token_input = nullptr; + // per token input + void* moe_reduction_token_input = nullptr; +}; + +template +struct MoeFinalizeAllReduceFusionParams : public AllReduceFusionParams { + // * moe reduction specific params + // Refer to kernel implementation on layout of those params + // number of active experts on current device + int top_k; + // [num_tokens, top_k] + void* expert_scale_factor = nullptr; + void* shared_expert_output = nullptr; + // [num_tokens, top_k] + int32_t* expanded_idx_to_permuted_idx = nullptr; + // allreduce_in [maxPermutedPaddedCount, hidden_dim] +}; + +template +struct LamportComm { + __device__ __forceinline__ LamportComm(void** workspace, int rank) { + counter_ptr = &reinterpret_cast(workspace[NRanks * 3])[0]; + flag_ptr = &reinterpret_cast(workspace[NRanks * 3])[2]; + clear_ptr = &reinterpret_cast(workspace[NRanks * 3])[4]; + flag_value = *flag_ptr; + int comm_size = reinterpret_cast(workspace[NRanks * 3])[3]; + clear_size = *clear_ptr; + int data_offset = flag_value % 3; + int clear_offset = (flag_value + 2) % 3; + for (int r = 0; r < NRanks; ++r) { + data_bufs[r] = reinterpret_cast(workspace[2 * NRanks + r]) + + static_cast(data_offset) * comm_size; + } + clear_buf = reinterpret_cast(workspace[2 * NRanks + rank]) + clear_offset * comm_size; + __syncthreads(); + if (threadIdx.x == 0) { + atomicAdd(counter_ptr, 1); + } + } + + __device__ __forceinline__ void update(int new_clear_size) { + if (blockIdx.x == 0 && threadIdx.x == 0) { + while (*reinterpret_cast(counter_ptr) != gridDim.x) { + } + *flag_ptr = (flag_value + 1) % 3; + *clear_ptr = new_clear_size; + *counter_ptr = 0; + } + } + + int* counter_ptr; + int* flag_ptr; + int* clear_ptr; + uint8_t* data_bufs[NRanks]; + uint8_t* clear_buf; + int clear_size; + int flag_value; +}; + +template +__device__ __forceinline__ vec_t vec_add(const vec_t& a, + const vec_t& b) { + vec_t ret; +#pragma unroll + for (int i = 0; i < VEC_SIZE; ++i) { + ret[i] = static_cast(a[i]) + static_cast(b[i]); + } + return ret; +} + +template +__device__ __forceinline__ vec_t rms_norm(vec_t const& residual, + vec_t const& gamma, + float const eps, int hidden_dim) { + __shared__ float s_val; + vec_t norm_out; + namespace cg = cooperative_groups; + cg::cluster_group cluster = cg::this_cluster(); + float acc = 0.f; +#pragma unroll + for (int i = 0; i < VEC_SIZE; ++i) { + float v = static_cast(residual[i]); + acc += v * v; + } + utils::blockReduceSumV2(&acc); + if (cluster.num_blocks() > 1) { + if (threadIdx.x == 0) { + s_val = acc; + acc = 0.f; + } + cluster.sync(); + if (threadIdx.x == 0) { + for (int i = 0; i < cluster.num_blocks(); ++i) { + acc += *cluster.map_shared_rank(&s_val, i); + } + } + cluster.sync(); + } + if (threadIdx.x == 0) { + s_val = rsqrtf(acc / hidden_dim + eps); + } + __syncthreads(); +#pragma unroll + for (int i = 0; i < VEC_SIZE; ++i) { + norm_out[i] = static_cast(residual[i]) * s_val * static_cast(gamma[i]); + } + return norm_out; +} + +template +__device__ __forceinline__ void fused_op(vec_t const& val, int access_id, int token_id, + int access_id_in_token, AllReduceFusionParams& params) { + if constexpr (AllReduceOut) { + val.store(reinterpret_cast(params.moe_allreduce_out) + access_id * VEC_SIZE); + } + vec_t residual_val; + residual_val.load(reinterpret_cast(params.residual_in) + access_id * VEC_SIZE); + + vec_t gamma_val; + gamma_val.load(reinterpret_cast(params.rms_gamma) + access_id_in_token * VEC_SIZE); + residual_val = vec_add(val, residual_val); + if constexpr (ResidualOut) { + residual_val.store(reinterpret_cast(params.residual_out) + access_id * VEC_SIZE); + } + vec_t norm_val; + norm_val = rms_norm(residual_val, gamma_val, params.rms_eps, params.hidden_dim); + if constexpr (NormOut) { + norm_val.store(reinterpret_cast(params.norm_out) + access_id * VEC_SIZE); + } + if constexpr (QuantOut) { + constexpr int SF_VEC_SIZE = 16; + auto sf_out = utils::cvt_quant_to_fp4_get_sf_out_offset( + std::nullopt /* batchIdx */, token_id, access_id_in_token, std::nullopt /* numRows */, + params.hidden_dim, reinterpret_cast(params.scale_out), params.layout); + reinterpret_cast(params.quant_out)[access_id] = + utils::cvt_warp_fp16_to_fp4(norm_val, params.scale_factor, sf_out); + } +} + +template +struct neg_zero { + static constexpr T value = -T(0); +}; + +template <> +struct neg_zero { + static constexpr unsigned short neg_zero_bits = 0x8000U; + static constexpr __half value = __half_raw{neg_zero_bits}; +}; + +template <> +struct neg_zero { + static constexpr unsigned short neg_zero_bits = 0x8000U; + static constexpr __nv_bfloat16 value = __nv_bfloat16_raw{neg_zero_bits}; +}; + +template <> +struct neg_zero { + static constexpr unsigned int neg_zero_bits = 0x80000000U; +}; + +template +__device__ static constexpr T neg_zero_v = neg_zero::value; + +template +__device__ bool is_negative_zero(T) { + return false; +} + +// float specialization +template <> +__device__ bool is_negative_zero(float x) { + return (__float_as_int(x) == 0x80000000); +} + +// double specialization +template <> +__device__ bool is_negative_zero(double x) { + return (__double_as_longlong(x) == 0x8000000000000000ULL); +} + +// __half specialization +template <> +__device__ bool is_negative_zero<__half>(__half x) { + return (__half_as_ushort(x) == 0x8000); +} + +// __nv_bfloat16 specialization +template <> +__device__ bool is_negative_zero<__nv_bfloat16>(__nv_bfloat16 x) { + return (__bfloat16_as_ushort(x) == 0x8000); +} + +template +__device__ __forceinline__ bool has_neg_zero(const vec_t& vec) { +#pragma unroll + for (int i = 0; i < VEC_SIZE; ++i) { + if (is_negative_zero(vec[i])) { + return true; + } + } + return false; +} + +template +__device__ __forceinline__ void remove_neg_zero(vec_t& vec) { +#pragma unroll + for (int i = 0; i < VEC_SIZE; ++i) { + vec[i] = (is_negative_zero(vec[i])) ? static_cast(0.f) : vec[i]; + } +} + +template +__device__ __forceinline__ void set_neg_zero(T* addr) { + vec_t val; + val.fill(neg_zero_v); + val.store_global_volatile(addr); +} + +int get_sm_count() { + static int sm_count = 0; + if (sm_count == 0) { + int device_id; + auto status = cudaGetDevice(&device_id); + FLASHINFER_CHECK(status == cudaSuccess, "cudaGetDevice failed with error code " + + std::string(cudaGetErrorString(status))); + cudaDeviceProp device_prop; + cudaGetDeviceProperties(&device_prop, device_id); + sm_count = device_prop.multiProcessorCount; + } + return sm_count; +} + +bool use_oneshot(int token_num) { return token_num <= details::kOneShotMaxToken; } + +///////////////////////////////////////////////////////////////// +// * MoE Reduction Fusion * // +///////////////////////////////////////////////////////////////// + +template +__global__ void moereduce_allreduce_fusion_kernel_oneshot_lamport( + MoeReductionAllReduceFusionParams params) { +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.wait;"); +#endif + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + namespace cg = cooperative_groups; + cg::cluster_group cluster = cg::this_cluster(); + cg::grid_group grid = cg::this_grid(); + + // Each token is handled by one cluster + // which token is handled by current cluster + int token_id = grid.cluster_rank(); + // total number of token + int num_token = params.size / params.hidden_dim; + // Each thread handle VEC_SIZE num elem in token. Total cluster.num_threads() to handle one + // token For current token, which VEC_SIZE is handled by current thread (in unit of + // VEC_SIZE) + int access_id_in_token = cluster.thread_rank(); + // Across all token, which VEC_SIZE is handled by current thread (in unit of + // VEC_SIZE) + static constexpr int VEC_SIZE = details::kBytesPerAccess / sizeof(T); + int access_id = token_id * params.hidden_dim / VEC_SIZE + access_id_in_token; + // Persistent kernel + // stride to next token handled by current cta + int token_stride = grid.num_clusters(); + // stride in unit of VEC_SIZE + int access_stride = token_stride * params.hidden_dim / VEC_SIZE; + // Total number of access in unit of VEC_SIZE to handle (token_num * hidden_dim) + // This is within one rank + int tot_access = params.size / VEC_SIZE; + vec_t clear_vec; + clear_vec.fill(neg_zero_v); + + cudaGridDependencySynchronize(); + LamportComm comm(params.workspace, params.rank); + int clear_access = comm.clear_size / VEC_SIZE; + + // * MoE related + int threadid_in_cluster = cluster.thread_rank(); + // Start Offset within one token's hidden_size of element + // Current thread handle token[thread_offset_within_token : thread_offset_within_token + + // VEC_SIZE] + int thread_offset_within_token = threadid_in_cluster * VEC_SIZE; + + // Persistent Kernel + // Each cluster iterate through all token it need to handle + for (int token_id = grid.cluster_rank(); token_id < num_token; token_id += grid.num_clusters()) { + if (thread_offset_within_token >= params.hidden_dim) { + break; + } + + // * MoE Reduce + // Offset within (num_token, hidden_size) in unit of element + int thread_offset_across_token = token_id * params.hidden_dim + thread_offset_within_token; + + vec_t accumulator; + accumulator.fill(0.f); + + // * Iterate through all active expert + int num_actexp = params.moe_reduction_device_num_experts; + for (int actexp_i = 0; actexp_i < num_actexp; ++actexp_i) { + // * Load active expert i's token j's partial data + // Offset within (num_act_exp, num_token, hidden_size) in unit of element + int thread_offset_across_actexp_token = + actexp_i * (params.hidden_dim * num_token) + thread_offset_across_token; + vec_t actexp_i_data; + actexp_i_data.load(reinterpret_cast(params.moe_reduction_active_experts_token_input) + + thread_offset_across_actexp_token); + + // * Load active expert i's token j's scale + int thread_offset_scale = actexp_i * num_token + token_id; + float actexp_i_token_j_scale = + reinterpret_cast(params.moe_reduction_scale_input)[thread_offset_scale]; + +#pragma unroll + for (int i = 0; i < VEC_SIZE; ++i) { + // assume computation is done in ScaleType + accumulator[i] += + static_cast((static_cast(actexp_i_data[i]) * actexp_i_token_j_scale)); + } + } + + // * FC2 + reduced(gGEMM2) + vec_t fc2_data; + fc2_data.load(reinterpret_cast(params.moe_reduction_token_input) + + thread_offset_across_token); + accumulator = vec_add(accumulator, fc2_data); + + // * AR Store + int access_id = token_id * params.hidden_dim / VEC_SIZE + access_id_in_token; + int idx = access_id; + + remove_neg_zero(accumulator); + +#pragma unroll + for (int r = 0; r < NRanks; ++r) { + // STG.128 to remote rank + int offset = (params.rank * tot_access + idx) * VEC_SIZE; + accumulator.store(reinterpret_cast(comm.data_bufs[r]) + offset); + } + } + + // * Clear previous buffer + for (int idx = access_id; idx < clear_access; idx += access_stride) { + int offset = idx * VEC_SIZE; + clear_vec.store(reinterpret_cast(comm.clear_buf) + offset); + } + + // * AR Load + Fusion + for (int idx = access_id, tidx = token_id; idx < tot_access; + idx += access_stride, tidx += token_stride) { + // * AR Load + vec_t vals[NRanks]; + bool done = false; + while (!done) { + // printf("Rank %d poll AR Load with flag %d\n", params.rank, *comm.flag_ptr); + done = true; +#pragma unroll + for (int r_i = 0; r_i < NRanks; ++r_i) { + int r = (r_i + params.rank) % NRanks; + // LDG.128 from local rank + vals[r].load_global_volatile(reinterpret_cast(comm.data_bufs[params.rank]) + + (r * tot_access + idx) * VEC_SIZE); + done &= !has_neg_zero(vals[r]); + } + } + + vec_t sum_val = vals[0]; +#pragma unroll + for (int r = 1; r < NRanks; ++r) { + sum_val = vec_add(sum_val, vals[r]); + } + + // * Fuse + fused_op(sum_val, idx, tidx, + access_id_in_token, params); + } + comm.update(params.size * NRanks); + cudaTriggerProgrammaticLaunchCompletion(); +#endif + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.launch_dependents;"); +#endif +} + +template +cudaError_t launch_oneshot_moereduce_lamport(MoeReductionAllReduceFusionParams const& params, + cudaLaunchConfig_t& cfg) { + FLASHINFER_CUDA_CALL(cudaLaunchKernelEx( + &cfg, + moereduce_allreduce_fusion_kernel_oneshot_lamport, + params)); + return cudaSuccess; +} + +template +cudaError_t moereduction_allreduce_fusion_kernel_launcher( + MoeReductionAllReduceFusionParams const& params, bool launch_with_pdl) { + int token_num = params.size / params.hidden_dim; + bool oneshot = use_oneshot(token_num); + // todo(yingyi): support token_num > oneshot max token in another kernel + if (oneshot == false) { + FLASHINFER_LOG_WARN("expect one shot but got %d tokens, expect performance degradation", + token_num); + oneshot = true; + } + // FLASHINFER_CHECK(oneshot, "only support one shot"); + // Each token is handled by one cluster + int cluster_num = token_num; + // Total number of threads (within one cluster) that's need to handle one token + // given that each thread handle kElemsPerAccess + int threads_per_token = params.hidden_dim * sizeof(T) / details::kBytesPerAccess; + // Total number of warp (within one cluster) that's need to handle one token + // given that each thread handle kElemsPerAccess + int warps_per_token = (threads_per_token + 31) / 32; + int cluster_size = 8; + while (warps_per_token % cluster_size != 0) { + cluster_size /= 2; + } + int block_size = warps_per_token / cluster_size * 32; + FLASHINFER_CHECK(block_size <= 1024 && cluster_size > 0, + "block_size <= 1024 && cluster_size > 0"); + int sm_count = get_sm_count(); + int grid_size = (std::min(sm_count, cluster_num * cluster_size) / cluster_size) * cluster_size; + cudaLaunchConfig_t cfg; + cudaLaunchAttribute attribute[2]; + cfg.gridDim = grid_size; + cfg.blockDim = block_size; + cfg.dynamicSmemBytes = 0; + cfg.stream = params.stream; + attribute[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attribute[0].val.programmaticStreamSerializationAllowed = launch_with_pdl ? 1 : 0; + attribute[1].id = cudaLaunchAttributeClusterDimension; + attribute[1].val.clusterDim.x = cluster_size; + attribute[1].val.clusterDim.y = 1; + attribute[1].val.clusterDim.z = 1; + cfg.attrs = attribute; + cfg.numAttrs = 2; + if (oneshot) { + FLASHINFER_CUDA_CALL( + (launch_oneshot_moereduce_lamport( + params, cfg))); + } + return cudaSuccess; +} + +#define DISPATCH_BOOL_(expr, const_expr, ...) \ + [&]() -> cudaError_t { \ + if (expr) { \ + constexpr bool const_expr = true; \ + return __VA_ARGS__(); \ + } else { \ + constexpr bool const_expr = false; \ + return __VA_ARGS__(); \ + } \ + }() + +#define _DISPATCH_MOEREDUCTION_CASE(n_ranks_val, N_RANKS_VAR, ar, res, rms, quant, AR, RES, RMS, \ + QUANT, ...) \ + case n_ranks_val: { \ + constexpr int N_RANKS_VAR = n_ranks_val; \ + return DISPATCH_BOOL_(ar, AR, [&]() -> cudaError_t { \ + return DISPATCH_BOOL_(res, RES, [&]() -> cudaError_t { \ + return DISPATCH_BOOL_(rms, RMS, [&]() -> cudaError_t { \ + return DISPATCH_BOOL_(quant, QUANT, [&]() -> cudaError_t { return __VA_ARGS__(); }); \ + }); \ + }); \ + }); \ + } + +#define DISPATCH_MOEREDUCTION(n_ranks, ar, res, rms, quant, N_RANKS, AR, RES, RMS, QUANT, ...) \ + [&]() -> cudaError_t { \ + switch (n_ranks) { \ + _DISPATCH_MOEREDUCTION_CASE(2, N_RANKS, ar, res, rms, quant, AR, RES, RMS, QUANT, \ + __VA_ARGS__) \ + _DISPATCH_MOEREDUCTION_CASE(4, N_RANKS, ar, res, rms, quant, AR, RES, RMS, QUANT, \ + __VA_ARGS__) \ + _DISPATCH_MOEREDUCTION_CASE(8, N_RANKS, ar, res, rms, quant, AR, RES, RMS, QUANT, \ + __VA_ARGS__) \ + _DISPATCH_MOEREDUCTION_CASE(16, N_RANKS, ar, res, rms, quant, AR, RES, RMS, QUANT, \ + __VA_ARGS__) \ + default: \ + FLASHINFER_CHECK(false, "Unsupported n_ranks"); \ + return cudaErrorNotSupported; \ + } \ + }() + +template +cudaError_t moereduction_allreduce_fusion_op(MoeReductionAllReduceFusionParams const& params, + bool launch_with_pdl) { + FLASHINFER_CHECK(params.residual_in && params.rms_gamma, "residual_in and rms_gamma must be set"); + FLASHINFER_CHECK(params.moe_reduction_scale_input && + params.moe_reduction_active_experts_token_input && + params.moe_reduction_token_input, + "moe_reduction_scale_input, moe_reduction_active_experts_token_input and " + "moe_reduction_token_input must be set"); + FLASHINFER_CHECK(params.size % params.hidden_dim == 0, "size must be a multiple of hidden_dim"); + FLASHINFER_CHECK(params.hidden_dim * sizeof(T) % details::kBytesPerAccess == 0, + "hidden_dim * sizeof(T) must be a multiple of kBytesPerAccess"); + FLASHINFER_CHECK( + params.moe_allreduce_out || params.residual_out || params.norm_out || params.quant_out, + "at least one of moe_allreduce_out, residual_out, norm_out, quant_out must be set"); + + // hidden_dim (d) = 7168 for dpsk moe, and hence 128 tokens as one-shot threshold + // AR outputs are optional, since we always have fused options followed. + // pattern1: AR+Residual+Add_RMS+Quant + // [m, d] bf16 allreduce_in, [m, d] bf16 residual_in + // [m, d] bf16 residual_out, [m, d] fp4 quant_out + + // pattern2: AR+Add_RMS + // [m, d] bf16 allreduce_in, [m, d] bf16 residual_in + // [m, d] bf16 norm_out + + // pattern3: AR+Add_RMS + // [m, d] bf16 allreduce_in, [m, d] bf16 residual_in + // [m, d] bf16 norm_out + + // pattern4: AR+Add_RMS + // [m, d] bf16 allreduce_in, [m, d] bf16 residual_in + // [m, d] bf16 residual_out, [m, d] bf16 norm_out, [m, d] fp4 quant_out + + auto status = DISPATCH_MOEREDUCTION( + params.nranks, params.moe_allreduce_out, params.residual_out, params.rms_gamma, + params.quant_out, N_RANKS, AR, RES, RMS, QUANT, [&]() -> cudaError_t { + FLASHINFER_CUDA_CALL( + (moereduction_allreduce_fusion_kernel_launcher( + (params), (launch_with_pdl)))); + }); + return status; +} + +///////////////////////////////////////////////////////////////// +// * MoE Finalize Allreduce Fusion * // +///////////////////////////////////////////////////////////////// + +template +__global__ void moefinalize_allreduce_fusion_kernel_oneshot_lamport( + MoeFinalizeAllReduceFusionParams params) { +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + namespace cg = cooperative_groups; + cg::cluster_group cluster = cg::this_cluster(); + cg::grid_group grid = cg::this_grid(); + + static constexpr int VEC_SIZE = details::kBytesPerAccess / sizeof(T); + + // Each token is handled by one cluster + // which token is handled by current cluster + int token_id = grid.cluster_rank(); + // total number of token + int num_token = params.size / params.hidden_dim; + // Each thread handle VEC_SIZE num elem in token. Total cluster.num_threads() to handle one + // token For current token, which VEC_SIZE is handled by current thread (in unit of + // VEC_SIZE) + int access_id_in_token = cluster.thread_rank(); + // Across all token, which VEC_SIZE is handled by current thread (in unit of + // VEC_SIZE) + int access_id = token_id * params.hidden_dim / VEC_SIZE + access_id_in_token; + // Persistent kernel + // stride to next token handled by current cta + int token_stride = grid.num_clusters(); + // stride in unit of VEC_SIZE + int access_stride = token_stride * params.hidden_dim / VEC_SIZE; + // Total number of access in unit of VEC_SIZE to handle (token_num * hidden_dim) + // This is within one rank + int tot_access = params.size / VEC_SIZE; + vec_t clear_vec; + clear_vec.fill(neg_zero_v); + + cudaGridDependencySynchronize(); + LamportComm comm(params.workspace, params.rank); + int clear_access = comm.clear_size / VEC_SIZE; + + // * MoE related + int threadid_in_cluster = cluster.thread_rank(); + // Start Offset within one token's hidden_size of element + // Current thread handle token[thread_offset_within_token : thread_offset_within_token + + // VEC_SIZE] + int thread_offset_within_token = threadid_in_cluster * VEC_SIZE; + + int top_k = params.top_k; + bool use_scale_factor = params.expert_scale_factor != nullptr; + + // Persistent Kernel + // Each cluster iterate through all token it need to handle + for (int token_id = grid.cluster_rank(); token_id < num_token; token_id += grid.num_clusters()) { + if (thread_offset_within_token >= params.hidden_dim) { + break; + } + + // * MoE finalize + vec_t accumulator; + accumulator.fill(0.f); + + for (int k = 0; k < top_k; k++) { + int const expanded_idx = token_id * top_k + k; + int32_t const permuted_idx = params.expanded_idx_to_permuted_idx[expanded_idx]; + + if (permuted_idx == -1) continue; + + int thread_offset_across_token = + permuted_idx * params.hidden_dim + thread_offset_within_token; + float block_scale = 1.0; + if (use_scale_factor) { + block_scale = + static_cast(static_cast(params.expert_scale_factor)[expanded_idx]); + } + + vec_t permuted_data; + permuted_data.load(reinterpret_cast(params.allreduce_in) + thread_offset_across_token); + + // * acc += scale(data) +#pragma unroll + for (int i = 0; i < VEC_SIZE; ++i) { + // assume computation is done in ScaleType + accumulator[i] += static_cast(static_cast(permuted_data[i]) * block_scale); + } + } + + // * Add shared expert output + if (params.shared_expert_output) { + // * Load shared expert output + int thread_offset_across_token = token_id * params.hidden_dim + thread_offset_within_token; + vec_t shared_expert_output; + shared_expert_output.load(reinterpret_cast(params.shared_expert_output) + + thread_offset_across_token); +#pragma unroll + accumulator = vec_add(accumulator, shared_expert_output); + } + + // * AR Store + int idx = token_id * params.hidden_dim / VEC_SIZE + access_id_in_token; + remove_neg_zero(accumulator); + +#pragma unroll + for (int r = 0; r < NRanks; ++r) { + // STG.128 to remote rank + int offset = (params.rank * tot_access + idx) * VEC_SIZE; + accumulator.store_global_volatile(reinterpret_cast(comm.data_bufs[r]) + offset); + } + } + + // * Clear previous buffer + for (int idx = access_id; idx < clear_access; idx += access_stride) { + clear_vec.store(reinterpret_cast(comm.clear_buf) + idx * VEC_SIZE); + } + + // * AR Load + Fusion + for (int idx = access_id, tidx = token_id; idx < tot_access; + idx += access_stride, tidx += token_stride) { + // * AR Load + vec_t vals[NRanks]; + bool done = false; + while (!done) { + done = true; +#pragma unroll + for (int r = 0; r < NRanks; ++r) { + // LDG.128 from local rank + vals[r].load_global_volatile(reinterpret_cast(comm.data_bufs[r]) + + (r * tot_access + idx) * VEC_SIZE); + done &= !has_neg_zero(vals[r]); + } + } + vec_t sum_val = vals[0]; +#pragma unroll + for (int r = 1; r < NRanks; ++r) { + sum_val = vec_add(sum_val, vals[r]); + } + + // * Fuse: AllReduceOut is always false in finalize_moe_allreduce + fused_op(sum_val, idx, tidx, + access_id_in_token, params); + } + comm.update(params.size * NRanks); + cudaTriggerProgrammaticLaunchCompletion(); +#endif +} + +template +cudaError_t launch_oneshot_moefinalize_lamport(MoeFinalizeAllReduceFusionParams const& params, + cudaLaunchConfig_t& cfg) { + FLASHINFER_CUDA_CALL(cudaLaunchKernelEx( + &cfg, + moefinalize_allreduce_fusion_kernel_oneshot_lamport, + params)); + return cudaSuccess; +} + +template +cudaError_t moefinalize_allreduce_fusion_kernel_launcher( + MoeFinalizeAllReduceFusionParams const& params, bool launch_with_pdl) { + int token_num = params.size / params.hidden_dim; + bool oneshot = use_oneshot(token_num); + if (oneshot == false) { + FLASHINFER_LOG_WARN("expect one shot but got %d tokens, expect performance degradation", + token_num); + oneshot = true; + } + // Only support one shot + // FLASHINFER_CHECK(oneshot, "only support one shot"); + // Each token is handled by one cluster + int cluster_num = token_num; + // Total number of threads (within one cluster) that's need to handle one token + // given that each thread handle VEC_SIZE + static constexpr int VEC_SIZE = details::kBytesPerAccess / sizeof(T); + int threads_per_token = params.hidden_dim / VEC_SIZE; + // Total number of warp (within one cluster) that's need to handle one token + // given that each thread handle VEC_SIZE + int warps_per_token = (threads_per_token + 31) / 32; + int cluster_size = 8; + while (warps_per_token % cluster_size != 0) { + cluster_size /= 2; + } + int block_size = warps_per_token / cluster_size * 32; + FLASHINFER_CHECK(block_size <= 1024 && cluster_size > 0, + "block_size <= 1024 && cluster_size > 0"); + int sm_count = get_sm_count(); + int grid_size = (std::min(sm_count, cluster_num * cluster_size) / cluster_size) * cluster_size; + cudaLaunchConfig_t cfg; + cudaLaunchAttribute attribute[2]; + cfg.gridDim = grid_size; + cfg.blockDim = block_size; + cfg.dynamicSmemBytes = 0; + cfg.stream = params.stream; + attribute[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attribute[0].val.programmaticStreamSerializationAllowed = launch_with_pdl ? 1 : 0; + attribute[1].id = cudaLaunchAttributeClusterDimension; + attribute[1].val.clusterDim.x = cluster_size; + attribute[1].val.clusterDim.y = 1; + attribute[1].val.clusterDim.z = 1; + cfg.attrs = attribute; + cfg.numAttrs = 2; + if (oneshot) { + FLASHINFER_CUDA_CALL( + (launch_oneshot_moefinalize_lamport( + params, cfg))); + } + return cudaSuccess; +} + +#define _DISPATCH_MOEFINALIZEREDUCTION_CASE(n_ranks_val, N_RANKS_VAR, res, rms, quant, RES, RMS, \ + QUANT, ...) \ + case n_ranks_val: { \ + constexpr int N_RANKS_VAR = n_ranks_val; \ + return DISPATCH_BOOL_(res, RES, [&]() -> cudaError_t { \ + return DISPATCH_BOOL_(rms, RMS, [&]() -> cudaError_t { \ + return DISPATCH_BOOL_(quant, QUANT, [&]() -> cudaError_t { return __VA_ARGS__(); }); \ + }); \ + }); \ + } + +#define DISPATCH_MOEFINALIZEREDUCTION(n_ranks, res, rms, quant, N_RANKS, RES, RMS, QUANT, ...) \ + [&]() -> cudaError_t { \ + switch (n_ranks) { \ + _DISPATCH_MOEFINALIZEREDUCTION_CASE(2, N_RANKS, res, rms, quant, RES, RMS, QUANT, \ + __VA_ARGS__) \ + _DISPATCH_MOEFINALIZEREDUCTION_CASE(4, N_RANKS, res, rms, quant, RES, RMS, QUANT, \ + __VA_ARGS__) \ + _DISPATCH_MOEFINALIZEREDUCTION_CASE(8, N_RANKS, res, rms, quant, RES, RMS, QUANT, \ + __VA_ARGS__) \ + _DISPATCH_MOEFINALIZEREDUCTION_CASE(16, N_RANKS, res, rms, quant, RES, RMS, QUANT, \ + __VA_ARGS__) \ + default: \ + FLASHINFER_CHECK(false, "Unsupported n_ranks"); \ + return cudaErrorNotSupported; \ + } \ + }() + +template +cudaError_t moefinalize_allreduce_fusion_op(MoeFinalizeAllReduceFusionParams const& params, + bool launch_with_pdl) { + static constexpr int VEC_SIZE = details::kBytesPerAccess / sizeof(T); + FLASHINFER_CHECK(params.allreduce_in && params.expanded_idx_to_permuted_idx && params.top_k, + "allreduce_in, expanded_idx_to_permuted_idx and top_k must be set"); + FLASHINFER_CHECK(params.size % params.hidden_dim == 0, "size must be a multiple of hidden_dim"); + FLASHINFER_CHECK(params.hidden_dim % VEC_SIZE == 0, "hidden_dim must be a multiple of VEC_SIZE"); + + auto status = DISPATCH_MOEFINALIZEREDUCTION( + params.nranks, params.residual_out, params.rms_gamma, params.quant_out, N_RANKS, RES, RMS, + QUANT, [&]() -> cudaError_t { + FLASHINFER_CUDA_CALL( + (moefinalize_allreduce_fusion_kernel_launcher( + (params), (launch_with_pdl)))); + }); + return status; +} +} // namespace trtllm_moe_allreduce_fusion +} // namespace flashinfer diff --git a/include/flashinfer/comm/vllm_custom_all_reduce.cuh b/include/flashinfer/comm/vllm_custom_all_reduce.cuh new file mode 100644 index 0000000000000000000000000000000000000000..976530177aa4790e45968e7e6aef3400dbf99da4 --- /dev/null +++ b/include/flashinfer/comm/vllm_custom_all_reduce.cuh @@ -0,0 +1,512 @@ +// flashinfer: adapted from sglang + vllm code +// refer to: https://github.com/vllm-project/vllm/blob/v0.8.2/csrc/custom_all_reduce.cuh +#pragma once + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +struct cuda_error : public std::runtime_error { + /** + * @brief Constructs a `cuda_error` object with the given `message`. + * + * @param message The error char array used to construct `cuda_error` + */ + cuda_error(const char* message) : std::runtime_error(message) {} + /** + * @brief Constructs a `cuda_error` object with the given `message` string. + * + * @param message The `std::string` used to construct `cuda_error` + */ + cuda_error(std::string const& message) : cuda_error{message.c_str()} {} +}; + +#define CHECK_CUDA_SUCCESS(cmd) \ + do { \ + cudaError_t e = cmd; \ + if (e != cudaSuccess) { \ + std::stringstream _message; \ + auto s = cudaGetErrorString(e); \ + _message << std::string(s) + "\n" << __FILE__ << ':' << __LINE__; \ + throw cuda_error(_message.str()); \ + } \ + } while (0) + +namespace vllm { + +constexpr int kMaxBlocks = 36; +// Counter may overflow, but it's fine since unsigned int overflow is +// well-defined behavior. +using FlagType = uint32_t; +struct Signal { + alignas(128) FlagType self_counter[kMaxBlocks][8]; + // Two sets of peer counters are needed for two syncs. The reason is that + // it's possible for peer GPU block to arrive at the second sync point while + // the current GPU block haven't passed the first sync point. Thus, peer GPU + // may write counter+1 while current GPU is busy waiting for counter. We use + // alternating counter array to avoid this possibility. + alignas(128) FlagType peer_counter[2][kMaxBlocks][8]; +}; + +struct __align__(16) RankData { + void* __restrict__ ptrs[8]; +}; + +struct __align__(16) RankSignals { + Signal* signals[8]; +}; + +// like std::array, but aligned +template +struct __align__(alignof(T) * sz) array_t { + T data[sz]; + using type = T; + static constexpr int size = sz; +}; + +// use packed type to maximize memory efficiency +// goal: generate ld.128 and st.128 instructions +template +struct packed_t { + // the (P)acked type for load/store + using P = array_t; + // the (A)ccumulator type for reduction + using A = array_t; +}; + +#define DINLINE __device__ __forceinline__ + +// scalar cast functions +DINLINE float upcast_s(half val) { return __half2float(val); } + +template +DINLINE T downcast_s(float val); +template <> +DINLINE half downcast_s(float val) { + return __float2half(val); +} + +// scalar add functions +// for some reason when compiling with Pytorch, the + operator for half and +// bfloat is disabled so we call the intrinsics directly +DINLINE half& assign_add(half& a, half b) { + a = __hadd(a, b); + return a; +} +DINLINE float& assign_add(float& a, float b) { return a += b; } + +#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) +DINLINE float upcast_s(nv_bfloat16 val) { return __bfloat162float(val); } +template <> +DINLINE nv_bfloat16 downcast_s(float val) { + return __float2bfloat16(val); +} +DINLINE nv_bfloat16& assign_add(nv_bfloat16& a, nv_bfloat16 b) { + a = __hadd(a, b); + return a; +} +#endif + +template +DINLINE array_t& packed_assign_add(array_t& a, array_t b) { +#pragma unroll + for (int i = 0; i < N; i++) { + assign_add(a.data[i], b.data[i]); + } + return a; +} + +template +DINLINE array_t upcast(array_t val) { + if constexpr (std::is_same::value) { + return val; + } else { + array_t out; +#pragma unroll + for (int i = 0; i < N; i++) { + out.data[i] = upcast_s(val.data[i]); + } + return out; + } +} + +template +DINLINE O downcast(array_t val) { + if constexpr (std::is_same::value) { + return val; + } else { + O out; +#pragma unroll + for (int i = 0; i < O::size; i++) { + out.data[i] = downcast_s(val.data[i]); + } + return out; + } +} + +static DINLINE void st_flag_release(FlagType* flag_addr, FlagType flag) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700 + asm volatile("st.release.sys.global.u32 [%1], %0;" ::"r"(flag), "l"(flag_addr)); +#else + asm volatile("membar.sys; st.volatile.global.u32 [%1], %0;" ::"r"(flag), "l"(flag_addr)); +#endif +} + +static DINLINE FlagType ld_flag_acquire(FlagType* flag_addr) { + FlagType flag; +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700 + asm volatile("ld.acquire.sys.global.u32 %0, [%1];" : "=r"(flag) : "l"(flag_addr)); +#else + asm volatile("ld.volatile.global.u32 %0, [%1]; membar.gl;" : "=r"(flag) : "l"(flag_addr)); +#endif + return flag; +} + +static DINLINE void st_flag_volatile(FlagType* flag_addr, FlagType flag) { + asm volatile("st.volatile.global.u32 [%1], %0;" ::"r"(flag), "l"(flag_addr)); +} + +static DINLINE FlagType ld_flag_volatile(FlagType* flag_addr) { + FlagType flag; + asm volatile("ld.volatile.global.u32 %0, [%1];" : "=r"(flag) : "l"(flag_addr)); + return flag; +} + +// is_start: whether this is the very first synchronization barrier. +// need_fence: whether a memory fence is needed. If true, a release-acquire +// semantic is used to enforce memory access order before and after this +// barrier. +template +DINLINE void multi_gpu_barrier(const RankSignals& sg, Signal* self_sg, int rank) { + if constexpr (!is_start) __syncthreads(); + static_assert(!(is_start && need_fence)); // Start barrier shouldn't need fence. + if (threadIdx.x < ngpus) { + // Increment the counter. Technically we only need one counter, but we use + // multiple per block to eliminate the need to share the counter via smem. + auto val = self_sg->self_counter[blockIdx.x][threadIdx.x] += 1; + // Write the expected counter value to peer and wait for correct value from + // peer. + auto peer_counter_ptr = &sg.signals[threadIdx.x]->peer_counter[val % 2][blockIdx.x][rank]; + auto self_counter_ptr = &self_sg->peer_counter[val % 2][blockIdx.x][threadIdx.x]; + if constexpr (need_fence) { + st_flag_release(peer_counter_ptr, val); + while (ld_flag_acquire(self_counter_ptr) != val); + } else { + st_flag_volatile(peer_counter_ptr, val); + while (ld_flag_volatile(self_counter_ptr) != val); + } + } + if constexpr (is_start || need_fence) __syncthreads(); +} + +template +DINLINE P packed_reduce(const P* ptrs[], int idx) { + A tmp = upcast(ptrs[0][idx]); +#pragma unroll + for (int i = 1; i < ngpus; i++) { + packed_assign_add(tmp, upcast(ptrs[i][idx])); + } + return downcast

(tmp); +} + +template +__global__ void __launch_bounds__(512, 1) + cross_device_reduce_1stage(RankData* _dp, RankSignals sg, Signal* self_sg, + T* __restrict__ result, int rank, int size) { + using P = typename packed_t::P; + using A = typename packed_t::A; + // note: we don't reorder the address so the accumulation order is the same + // for all ranks, ensuring bitwise identical results + auto dp = *_dp; + multi_gpu_barrier(sg, self_sg, rank); + // do the actual reduction + for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < size; idx += gridDim.x * blockDim.x) { + ((P*)result)[idx] = packed_reduce((const P**)&dp.ptrs[0], idx); + } + multi_gpu_barrier(sg, self_sg, rank); +} + +template +DINLINE P* get_tmp_buf(Signal* sg) { + return (P*)(((Signal*)sg) + 1); +} + +template +__global__ void __launch_bounds__(512, 1) + cross_device_reduce_2stage(RankData* _dp, RankSignals sg, Signal* self_sg, + T* __restrict__ result, int rank, int size) { + int tid = blockIdx.x * blockDim.x + threadIdx.x; + int stride = gridDim.x * blockDim.x; + using P = typename packed_t::P; + using A = typename packed_t::A; + int part = size / ngpus; + int start = rank * part; + int end = rank == ngpus - 1 ? size : start + part; + int largest_part = part + size % ngpus; + const P* ptrs[ngpus]; + P* tmps[ngpus]; +#pragma unroll + for (int i = 0; i < ngpus; i++) { + int target = (rank + i) % ngpus; + ptrs[i] = (const P*)_dp->ptrs[target]; + tmps[i] = get_tmp_buf

(sg.signals[target]); + } + auto tmp_out = tmps[0]; + multi_gpu_barrier(sg, self_sg, rank); + // stage 1: reduce scatter + for (int idx = start + tid; idx < end; idx += stride) { + tmp_out[idx - start] = packed_reduce(ptrs, idx); + } + multi_gpu_barrier(sg, self_sg, rank); + + // stage 2: allgather. Note: it's important to match the tid between + // the two stages, because visibility across devices is only guaranteed + // between threads that have the same tid. If thread i computes the sum of + // start + i in the first stage, then thread i also gathers start + i from all + // ranks. + for (int idx = tid; idx < largest_part; idx += stride) { +#pragma unroll + for (int i = 0; i < ngpus; i++) { + int gather_from_rank = ((rank + i) % ngpus); + if (gather_from_rank == ngpus - 1 || idx < part) { + int dst_idx = gather_from_rank * part + idx; + ((P*)result)[dst_idx] = tmps[i][idx]; + } + } + } +} + +using IPC_KEY = std::array; +static_assert(sizeof(IPC_KEY) == sizeof(cudaIpcMemHandle_t)); +static_assert(alignof(IPC_KEY) == alignof(cudaIpcMemHandle_t)); + +class CustomAllreduce { + public: + int rank_; + int world_size_; + bool full_nvlink_; + + RankSignals sg_; + // Stores an map from a pointer to its peer pointters from all ranks. + std::unordered_map buffers_; + Signal* self_sg_; + + // Stores rank data from all ranks. This is mainly for cuda graph purposes. + // For cuda graph to work, all kernel arguments must be fixed during graph + // capture time. However, the peer pointers are not known during graph capture + // time. Therefore, during capture, we increment the rank data pointer and use + // that as the argument to the kernel. The kernel arguments are stored in + // graph_unreg_buffers_. The actual peer pointers will be filled in at the + // memory pointed to by the pointers in graph_unreg_buffers_ when + // the IPC handles are exchanged between ranks. + // + // The overall process looks like this: + // 1. Graph capture. + // 2. Each rank obtains the IPC handles for each addresses used during cuda + // graph capture using get_graph_buffer_ipc_meta. + // 3. (In Python) all gather the IPC handles. + // 4. Obtain the peer pointers by opening the IPC handles, and store them in + // the rank data array at corresponding positions. + RankData *d_rank_data_base_, *d_rank_data_end_; + std::vector graph_unreg_buffers_; + // a map from IPC handles to opened IPC pointers + std::map ipc_handles_; + + /** + * Signals are an array of ipc-enabled buffers from all ranks. + * For each of the buffer, the layout is as follows: + * | -- sizeof(Signal) -- | ------ a few MB ----- | + * The first section is for allreduce synchronization, and the second section + * is for storing the intermediate results required by some allreduce algos. + * + * Note: this class does not own any device memory. Any required buffers + * are passed in from the constructor. + */ + CustomAllreduce(Signal** signals, void* rank_data, size_t rank_data_sz, int rank, int world_size, + bool full_nvlink = true) + : rank_(rank), + world_size_(world_size), + full_nvlink_(full_nvlink), + self_sg_(signals[rank]), + d_rank_data_base_(reinterpret_cast(rank_data)), + d_rank_data_end_(d_rank_data_base_ + rank_data_sz / sizeof(RankData)) { + for (int i = 0; i < world_size_; i++) { + sg_.signals[i] = signals[i]; + } + } + + char* open_ipc_handle(const void* ipc_handle) { + auto [it, new_handle] = ipc_handles_.insert({*((IPC_KEY*)ipc_handle), nullptr}); + if (new_handle) { + char* ipc_ptr; + CHECK_CUDA_SUCCESS(cudaIpcOpenMemHandle((void**)&ipc_ptr, + *((const cudaIpcMemHandle_t*)ipc_handle), + cudaIpcMemLazyEnablePeerAccess)); + it->second = ipc_ptr; + } + return it->second; + } + + std::pair> get_graph_buffer_ipc_meta() { + auto num_buffers = graph_unreg_buffers_.size(); + auto handle_sz = sizeof(cudaIpcMemHandle_t); + std::string handles(handle_sz * num_buffers, static_cast(0)); + std::vector offsets(num_buffers); + for (int i = 0; i < num_buffers; i++) { + auto ptr = graph_unreg_buffers_[i]; + void* base_ptr; + // note: must share the base address of each allocation, or we get wrong + // address + CHECK_CUDA_SUCCESS( + cudaIpcGetMemHandle((cudaIpcMemHandle_t*)&handles[i * handle_sz], base_ptr)); + offsets[i] = ((char*)ptr) - ((char*)base_ptr); + } + return std::make_pair(handles, offsets); + } + + void check_rank_data_capacity(size_t num = 1) { + if (d_rank_data_base_ + num > d_rank_data_end_) + throw std::runtime_error("Rank data buffer is overflowed by " + + std::to_string(d_rank_data_base_ + num - d_rank_data_end_)); + } + + /** + * Register already-shared IPC pointers. + */ + void register_buffer(void** ptrs) { + check_rank_data_capacity(); + RankData data; + for (int i = 0; i < world_size_; i++) { + data.ptrs[i] = ptrs[i]; + } + auto d_data = d_rank_data_base_++; + CHECK_CUDA_SUCCESS(cudaMemcpy(d_data, &data, sizeof(RankData), cudaMemcpyHostToDevice)); + buffers_[ptrs[rank_]] = d_data; + } + + // Note: when registering graph buffers, we intentionally choose to not + // deduplicate the addresses. That means if the allocator reuses some + // addresses, they will be registered again. This is to account for the remote + // possibility of different allocation patterns between ranks. For example, + // rank 1 may get the same input address for the second allreduce, but rank 2 + // got a different address. IPC handles have internal reference counting + // mechanism so overhead should be small. + void register_graph_buffers(const std::vector& handles, + const std::vector>& offsets) { + auto num_buffers = graph_unreg_buffers_.size(); + check_rank_data_capacity(num_buffers); + std::vector rank_data(num_buffers); + for (int i = 0; i < num_buffers; i++) { + auto self_ptr = graph_unreg_buffers_[i]; + auto& rd = rank_data[i]; + for (int j = 0; j < world_size_; j++) { + if (j != rank_) { + char* handle = open_ipc_handle(&handles[j][i * sizeof(cudaIpcMemHandle_t)]); + handle += offsets[j][i]; + rd.ptrs[j] = handle; + } else { + rd.ptrs[j] = self_ptr; + } + } + } + CHECK_CUDA_SUCCESS(cudaMemcpy(d_rank_data_base_, rank_data.data(), + sizeof(RankData) * num_buffers, cudaMemcpyHostToDevice)); + d_rank_data_base_ += num_buffers; + graph_unreg_buffers_.clear(); + } + + /** + * Performs allreduce, assuming input has already been registered. + * + * Block and grid default configs are results after careful grid search. Using + * 36 blocks give the best or close to the best runtime on the devices I + * tried: A100, A10, A30, T4, V100. You'll notice that NCCL kernels also only + * take a small amount of SMs. Not quite sure the underlying reason, but my + * guess is that too many SMs will cause contention on NVLink bus. + */ + template + void allreduce(cudaStream_t stream, T* input, T* output, int size, int block_limit, + int threads = 512) { + auto d = packed_t::P::size; + if (size % d != 0) + throw std::runtime_error( + "custom allreduce currently requires input length to be multiple " + "of " + + std::to_string(d)); + if (block_limit > kMaxBlocks) + throw std::runtime_error("max supported block limit is " + std::to_string(kMaxBlocks) + + ". Got " + std::to_string(block_limit)); + + RankData* ptrs; + cudaStreamCaptureStatus status; + CHECK_CUDA_SUCCESS(cudaStreamIsCapturing(stream, &status)); + if (status == cudaStreamCaptureStatusActive) { + ptrs = d_rank_data_base_ + graph_unreg_buffers_.size(); + graph_unreg_buffers_.push_back(input); + } else { + auto it = buffers_.find(input); + if (it == buffers_.end()) + throw std::runtime_error("buffer address " + + std::to_string(reinterpret_cast(input)) + + " is not registered!"); + ptrs = it->second; + } + + size /= d; + auto bytes = size * sizeof(typename packed_t::P); + int blocks = std::min(block_limit, (size + threads - 1) / threads); +#define KL(ngpus, name) \ + name<<>>(ptrs, sg_, self_sg_, output, rank_, size); + // TODO(hanzhi713): Threshold is different for A100 and H100. + // Add per device threshold. +#define REDUCE_CASE(ngpus) \ + case ngpus: { \ + if (world_size_ == 2) { \ + KL(ngpus, cross_device_reduce_1stage); \ + } else if (full_nvlink_) { \ + if ((world_size_ <= 4 && bytes < 512 * 1024) || (world_size_ <= 8 && bytes < 256 * 1024)) { \ + KL(ngpus, cross_device_reduce_1stage); \ + } else { \ + KL(ngpus, cross_device_reduce_2stage); \ + } \ + } \ + break; \ + } + + switch (world_size_) { + REDUCE_CASE(2) + REDUCE_CASE(4) + REDUCE_CASE(6) + REDUCE_CASE(8) + default: + throw std::runtime_error( + "custom allreduce only supports num gpus in (2,4,6,8). Actual num " + "gpus = " + + std::to_string(world_size_)); + } +#undef REDUCE_CASE +#undef KL + } + + ~CustomAllreduce() { + for (auto [_, ptr] : ipc_handles_) { + CHECK_CUDA_SUCCESS(cudaIpcCloseMemHandle(ptr)); + } + } +}; +/** + * To inspect PTX/SASS, copy paste this header file to compiler explorer and add + a template instantiation: + * template void vllm::CustomAllreduce::allreduce(cudaStream_t, half *, + half *, int, int, int); +*/ +} // namespace vllm diff --git a/include/flashinfer/cp_async.cuh b/include/flashinfer/cp_async.cuh new file mode 100644 index 0000000000000000000000000000000000000000..bd59cc58e38319e327998d35176665798f6e6afa --- /dev/null +++ b/include/flashinfer/cp_async.cuh @@ -0,0 +1,189 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef FLASHINFER_CP_ASYNC_CUH_ +#define FLASHINFER_CP_ASYNC_CUH_ + +#include + +#include + +namespace flashinfer { + +namespace cp_async { + +enum class SharedMemFillMode { + kFillZero, // Fill zero to shared memory when predicate is false + kNoFill // Do not fill zero to shared memory when predicate is false +}; + +enum class PrefetchMode { + kNoPrefetch, // Do not fetch additional data from global memory to L2 + kPrefetch // Fetch additional data from global memory to L2 +}; + +#if (__CUDACC_VER_MAJOR__ >= 11) +#if (!defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 800)) +#define FLASHINFER_CP_ASYNC_ENABLED +#endif +#endif + +/*! + * \brief Wrapper of PTX cp.async.commit_group instruction, commit all prior uncommitted + * cp.async instructions to a group + */ +__device__ __forceinline__ void commit_group() { +#ifdef FLASHINFER_CP_ASYNC_ENABLED + asm volatile("cp.async.commit_group;\n" ::); +#endif +} + +/*! + * \brief Wrapper of PTX cp.async.wait_group instruction + * \tparam n Wait till most recent n groups are committed + */ +template +__device__ __forceinline__ void wait_group() { +#ifdef FLASHINFER_CP_ASYNC_ENABLED + asm volatile("cp.async.wait_group %0;\n" ::"n"(n)); +#endif +} + +/*! + * \brief Wrapper of PTX cp.async.cg.shared.global instruction, asynchronously copy data from + * global memory to shared memory + * \tparam prefetch_mode Whether to fetch additional data from global memory to L2 + * \tparam T Data type + * \param smem_ptr Pointer to shared memory + * \param gmem_ptr Pointer to global memory + */ +template +__device__ __forceinline__ void load_128b(T* smem_ptr, const T* gmem_ptr) { +#ifdef FLASHINFER_CP_ASYNC_ENABLED + uint32_t smem_int_ptr = static_cast(__cvta_generic_to_shared(smem_ptr)); + if constexpr (prefetch_mode == PrefetchMode::kPrefetch) { + asm volatile("cp.async.cg.shared.global.L2::128B [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr), + "l"(gmem_ptr), "n"(16), "r"(16)); + } else { + asm volatile("cp.async.cg.shared.global [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr), + "l"(gmem_ptr), "n"(16), "r"(16)); + } +#else + *((uint4*)smem_ptr) = *((uint4*)gmem_ptr); +#endif +} + +/*! + * \brief Wrapper of PTX cp.async.cg.shared.global instruction, asynchronously copy data from + * global memory to shared memory with predicate. + * \tparam prefetch_mode Whether to fetch additional data from global memory to L2 + * \tparam fill_mode Whether to fill zero to shared memory when predicate is false + * \tparam T Data type + * \param smem_ptr Pointer to shared memory + * \param gmem_ptr Pointer to global memory + * \param predicate Predicate value + * \note fill zero is slower than not fill zero + */ +template +__device__ __forceinline__ void pred_load_128b(T* smem_ptr, const T* gmem_ptr, bool predicate) { +#ifdef FLASHINFER_CP_ASYNC_ENABLED + uint32_t smem_int_ptr = static_cast(__cvta_generic_to_shared(smem_ptr)); + if constexpr (fill_mode == SharedMemFillMode::kFillZero) { + int src_in_bytes = predicate ? 16 : 0; + if constexpr (prefetch_mode == PrefetchMode::kPrefetch) { + asm volatile("cp.async.cg.shared.global.L2::128B [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr), + "l"(gmem_ptr), "n"(16), "r"(src_in_bytes)); + } else { + asm volatile("cp.async.cg.shared.global [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr), + "l"(gmem_ptr), "n"(16), "r"(src_in_bytes)); + } + } else { + if constexpr (prefetch_mode == PrefetchMode::kPrefetch) { + asm volatile( + "{\n" + " .reg .pred p;\n" + " setp.ne.b32 p, %0, 0;\n" + " @p cp.async.cg.shared.global.L2::128B [%1], [%2], %3;\n" + "}\n" ::"r"((int)predicate), + "r"(smem_int_ptr), "l"(gmem_ptr), "n"(16)); + } else { + asm volatile( + "{\n" + " .reg .pred p;\n" + " setp.ne.b32 p, %0, 0;\n" + " @p cp.async.cg.shared.global [%1], [%2], %3;\n" + "}\n" ::"r"((int)predicate), + "r"(smem_int_ptr), "l"(gmem_ptr), "n"(16)); + } + } +#else + if (predicate) { + *((uint4*)smem_ptr) = *((uint4*)gmem_ptr); + } else { + if constexpr (fill_mode == SharedMemFillMode::kFillZero) { + *((uint4*)smem_ptr) = make_uint4(0, 0, 0, 0); + } + } +#endif +} + +/*! + * \brief Load specified number of bits per thread from global memory to shared memory + * \tparam num_bits Number of bits to load, must be 128 or 256 + * \tparam prefetch_mode Whether to fetch additional data from global memory to L2 + * \tparam T Data type + * \param smem_ptr Pointer to shared memory + * \param gmem_ptr Pointer to global memory + */ +template +__device__ __forceinline__ void load(T* smem_ptr, const T* gmem_ptr) { + static_assert(num_bits == 128 || num_bits == 256, "num_bits must be 128 or 256"); + if constexpr (num_bits == 128) { + load_128b(smem_ptr, gmem_ptr); + } else { + load_128b(smem_ptr, gmem_ptr); + load_128b(smem_ptr + 16 / sizeof(T), gmem_ptr + 16 / sizeof(T)); + } +} + +/*! + * \brief Load specified number of bits per thread from global memory to shared memory with + * predicate + * \tparam num_bits Number of bits to load, must be 128 or 256 + * \tparam prefetch_mode Whether to fetch additional data from global memory to L2 + * \tparam fill_mode Whether to fill zero to shared memory when predicate is false + * \tparam T Data type + * \param smem_ptr Pointer to shared memory + * \param gmem_ptr Pointer to global memory + * \param predicate Predicate value + * \note fill zero is slower than not fill zero + */ +template +__device__ __forceinline__ void pred_load(T* smem_ptr, const T* gmem_ptr, bool predicate) { + static_assert(num_bits == 128 || num_bits == 256, "num_bits must be 128 or 256"); + if constexpr (num_bits == 128) { + pred_load_128b(smem_ptr, gmem_ptr, predicate); + } else { + pred_load_128b(smem_ptr, gmem_ptr, predicate); + pred_load_128b(smem_ptr + 16 / sizeof(T), gmem_ptr + 16 / sizeof(T), + predicate); + } +} + +} // namespace cp_async + +} // namespace flashinfer + +#endif // FLASHINFER_CP_ASYNC_CUH_ diff --git a/include/flashinfer/cubin_loader.h b/include/flashinfer/cubin_loader.h new file mode 100644 index 0000000000000000000000000000000000000000..7a5552a70146cac5472434781fc08c21c7e9b981 --- /dev/null +++ b/include/flashinfer/cubin_loader.h @@ -0,0 +1,86 @@ +/* + * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// This file exposes a C API to manage dynamic cubin loading. +// It is meant to be included in a .so file that is loaded dynamically +// as a python module, where the following functions can be referenced. +// +// The python code is expected to call `FlashInferSetCubinCallback` when +// loading the module. This sets a callback that can be called by the +// library to load a cubin. +// The callback is expected to call `FlashInferSetCurrentCubin` with the +// cubin data and the size of the cubin. +// +// Internally the library can just rely on the `getCubin` function to encapsulate +// this back and forth. +// +// This is a C API so that we can use it with ctypes and don't rely on pybind11, +// because pybind11 support of arbitrary callback requires >=python3.11. + +// Callback into the python function that will get us the requested cubin. +void (*callbackGetCubin)(const char* path, const char* sha256) = nullptr; + +// Set the python callback, called by the python code using ctypes. +extern "C" void FlashInferSetCubinCallback(void (*callback)(const char* path, const char* sha256)) { + callbackGetCubin = callback; +} + +// Thread-local variable that stores the current cubin. +// It is reset on every call to `getCubin()`. +thread_local std::string current_cubin; + +// Called by the callback to set the current cubin. +extern "C" void FlashInferSetCurrentCubin(const char* binary, int size) { + current_cubin = std::string(binary, size); +} + +// Get the cubin from the python callback. +// This is the API for the native library to use. +std::string getCubin(const std::string& name, const std::string& sha256) { + if (!callbackGetCubin) { + throw std::runtime_error("FlashInferSetCubinCallback not set"); + } + callbackGetCubin(name.c_str(), sha256.c_str()); + return current_cubin; +} + +void (*callbackGetMetaInfo)(const char* path, const char* sha256, const char* extension) = nullptr; + +// Set the python callback, called by the python code using ctypes. +extern "C" void FlashInferSetMetaInfoCallback(void (*callback)(const char* path, const char* sha256, + const char* extension)) { + callbackGetMetaInfo = callback; +} + +// Thread-local variable that stores the current metainfo. +// It is reset on every call to `getMetaInfo()`. +thread_local std::string raw_metainfo; + +// Called by the callback to set the current metainfo. +extern "C" void FlashInferSetCurrentMetaInfo(const char* binary, int size) { + raw_metainfo = std::string(binary, size); +} + +// Get the metainfo from the python callback. +// This is the API for the native library to use. +std::string getMetaInfo(const std::string& name, const std::string& sha256, + const std::string& extension) { + if (!callbackGetMetaInfo) { + throw std::runtime_error("FlashInferSetMetaInfoCallback not set"); + } + callbackGetMetaInfo(name.c_str(), sha256.c_str(), extension.c_str()); + return raw_metainfo; +} diff --git a/include/flashinfer/cutlass_utils.cuh b/include/flashinfer/cutlass_utils.cuh new file mode 100644 index 0000000000000000000000000000000000000000..4a101f434c3ee35a8b567f16f4bafefb00137d41 --- /dev/null +++ b/include/flashinfer/cutlass_utils.cuh @@ -0,0 +1,109 @@ +/* + * Copyright (c) 2024 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef FLASHINFER_CUTLASS_UTILS_CUH_ +#define FLASHINFER_CUTLASS_UTILS_CUH_ + +#include + +#include "cute/tensor.hpp" +#include "cutlass/cutlass.h" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/epilogue/collective/default_epilogue.hpp" +#include "cutlass/epilogue/thread/linear_combination.h" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/gemm/device/gemm_grouped.h" +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/dispatch_policy.hpp" +#include "cutlass/gemm/group_array_problem_shape.hpp" +#include "cutlass/gemm/kernel/default_gemm_grouped.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" +#include "cutlass/layout/matrix.h" +#include "cutlass/numeric_types.h" +#include "cutlass/tensor_ref.h" +#include "cutlass/util/command_line.h" +#include "cutlass/util/distribution.h" +#include "cutlass/util/host_tensor.h" +#include "cutlass/util/packed_stride.hpp" +#include "cutlass/util/reference/device/gemm.h" +#include "cutlass/util/reference/device/tensor_compare.h" +#include "cutlass/util/reference/device/tensor_fill.h" +#include "cutlass/util/tensor_view_io.h" +#if defined(FLASHINFER_ENABLE_FP4_E2M1) +#include +#endif + +namespace flashinfer { + +template +struct cutlass_dtype { + using type = T; +}; + +template <> +struct cutlass_dtype { + using type = cutlass::half_t; +}; + +template <> +struct cutlass_dtype { + using type = cutlass::bfloat16_t; +}; + +template <> +struct cutlass_dtype<__nv_fp8_e4m3> { + using type = cutlass::float_e4m3_t; +}; + +template <> +struct cutlass_dtype<__nv_fp8_e5m2> { + using type = cutlass::float_e5m2_t; +}; + +#if (__CUDACC_VER_MAJOR__ * 10000 + __CUDACC_VER_MINOR__ * 100 >= 120800) +template <> +struct cutlass_dtype<__nv_fp8_e8m0> { + using type = cutlass::float_ue8m0_t; +}; + +#if defined(FLASHINFER_ENABLE_FP4_E2M1) +template <> +struct cutlass_dtype<__nv_fp4_e2m1> { + using type = cutlass::float_e2m1_t; +}; +#endif +#endif + +template +using cutlass_dtype_t = typename cutlass_dtype::type; + +template +void compileTimeDebug(T&&) { + static_assert(sizeof(T) == 0, "Compile time debug"); +} + +#define CUTLASS_CHECK(cmd) \ + do { \ + auto status = cmd; \ + if (status != cutlass::Status::kSuccess) { \ + std::ostringstream err_msg; \ + err_msg << "cutlass " << #cmd << " failed: " << cutlassGetStatusString(status); \ + FLASHINFER_ERROR(err_msg.str()); \ + } \ + } while (0) + +} // namespace flashinfer + +#endif // FLASHINFER_CUTLASS_UTILS_CUH_ diff --git a/include/flashinfer/exception.h b/include/flashinfer/exception.h new file mode 100644 index 0000000000000000000000000000000000000000..1159511da4ba895f971e16e41bd718c75790a2b7 --- /dev/null +++ b/include/flashinfer/exception.h @@ -0,0 +1,61 @@ +/* + * Copyright (c) 2024 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef FLASHINFER_EXCEPTION_H_ +#define FLASHINFER_EXCEPTION_H_ + +#include +#include + +#define FLASHINFER_ERROR(message) throw flashinfer::Error(__FUNCTION__, __FILE__, __LINE__, message) + +template +void write_to_stream(std::ostringstream& oss, T&& val) { + oss << std::forward(val); +} + +template +void write_to_stream(std::ostringstream& oss, T&& val, Args&&... args) { + oss << std::forward(val) << " "; + write_to_stream(oss, std::forward(args)...); +} + +#define FLASHINFER_CHECK(condition, ...) \ + if (!(condition)) { \ + std::ostringstream oss; \ + write_to_stream(oss, __VA_ARGS__); \ + std::cerr << oss.str() << std::endl; \ + FLASHINFER_ERROR(oss.str()); \ + } + +namespace flashinfer { +class Error : public std::exception { + private: + std::string message_; + + public: + Error(const std::string& func, const std::string& file, int line, const std::string& message) { + std::ostringstream oss; + oss << "Error in function '" << func << "' " + << "at " << file << ":" << line << ": " << message; + message_ = oss.str(); + } + + virtual const char* what() const noexcept override { return message_.c_str(); } +}; + +} // namespace flashinfer + +#endif // FLASHINFER_EXCEPTION_H_ diff --git a/include/flashinfer/fastdiv.cuh b/include/flashinfer/fastdiv.cuh new file mode 100644 index 0000000000000000000000000000000000000000..305241c1af347caaceb2d02a9fbb779469c51579 --- /dev/null +++ b/include/flashinfer/fastdiv.cuh @@ -0,0 +1,111 @@ +/* + * Copyright 2014 Maxim Milakov + * + * The code is based on the Chapter 10 of Hacker's Delight book by Henry S. Warren, Jr. + * The struct is adapted from https://github.com/milakov/int_fastdiv/blob/master/int_fastdiv.h + * by Maxim Milakov, the difference is that here we use uint32_t instead of int32_t. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef FLASHINFER_FASTDIV_CUH_ +#define FLASHINFER_FASTDIV_CUH_ +#include + +namespace flashinfer { + +struct uint_fastdiv { + uint32_t d; + uint32_t m; + uint32_t s; + uint32_t a; + + __host__ __device__ uint_fastdiv() : d(0), m(0), s(0), a(0) {} + + __host__ uint_fastdiv(uint32_t d) : d(d) { + unsigned int p, nc, delta, q1, r1, q2, r2; + a = 0; + nc = unsigned(-1) - unsigned(-d) % d; + p = 31; + q1 = 0x80000000 / nc; + r1 = 0x80000000 - q1 * nc; + q2 = 0x7FFFFFFF / d; + r2 = 0x7FFFFFFF - q2 * d; + do { + p++; + if (r1 >= nc - r1) { + q1 = 2 * q1 + 1; + r1 = 2 * r1 - nc; + } else { + q1 = 2 * q1; + r1 = 2 * r1; + } + if (r2 + 1 >= d - r2) { + if (q2 >= 0x7FFFFFFF) a = 1; + q2 = 2 * q2 + 1; + r2 = 2 * r2 + 1 - d; + } else { + if (q2 >= 0x80000000) a = 1; + q2 = 2 * q2; + r2 = 2 * r2 + 1; + } + delta = d - 1 - r2; + } while (p < 64 && (q1 < delta || (q1 == delta && r1 == 0))); + m = q2 + 1; + s = p - 32; + } + + __host__ __device__ __forceinline__ operator unsigned int() const { return d; } + + __host__ __device__ __forceinline__ void divmod(uint32_t n, uint32_t& q, uint32_t& r) const { + if (d == 1) { + q = n; + } else { +#ifdef __CUDA_ARCH__ + q = __umulhi(m, n); +#else + q = (((unsigned long long)((long long)m * (long long)n)) >> 32); +#endif + q += a * n; + q >>= s; + } + r = n - q * d; + } +}; + +__host__ __device__ __forceinline__ uint32_t operator/(const uint32_t n, + const uint_fastdiv& divisor) { + uint32_t q; + if (divisor.d == 1) { + q = n; + } else { +#ifdef __CUDA_ARCH__ + q = __umulhi(divisor.m, n); +#else + q = (((unsigned long long)((long long)divisor.m * (long long)n)) >> 32); +#endif + q += divisor.a * n; + q >>= divisor.s; + } + return q; +} + +__host__ __device__ __forceinline__ uint32_t operator%(const uint32_t n, + const uint_fastdiv& divisor) { + uint32_t quotient = n / divisor; + uint32_t remainder = n - quotient * divisor; + return remainder; +} + +} // namespace flashinfer + +#endif // FLASHINFER_FASTDIV_CUH_ diff --git a/include/flashinfer/fp16.h b/include/flashinfer/fp16.h new file mode 100644 index 0000000000000000000000000000000000000000..d45d07817269660a43fbe5f9afcc236a8a1529b3 --- /dev/null +++ b/include/flashinfer/fp16.h @@ -0,0 +1,177 @@ +// SPDX - FileCopyrightText : 2017 - 2024 Marat Dukhan +// SPDX - FileCopyrightText : 2025 Advanced Micro Devices, Inc. +// +// SPDX - License - Identifier : MIT + +#pragma once + +#ifndef FLASHINFER_FP16_H +#define FLASHINFER_FP16_H + +#include +#include +#include +#include + +/* + * Convert a 32-bit floating-point number in IEEE single-precision format to a + * 16-bit floating-point number in IEEE half-precision format, in bit + * representation. + * + * @note The implementation relies on IEEE-like (no assumption about rounding + * mode and no operations on denormals) floating-point operations and bitcasts + * between integer and floating-point variables. + */ +static constexpr uint16_t fp16_ieee_from_fp32_value(float f) { + const float scale_to_inf = std::bit_cast(UINT32_C(0x77800000)); + const float scale_to_zero = std::bit_cast(UINT32_C(0x08800000)); + const float saturated_f = boost::math::ccmath::fabs(f) * scale_to_inf; + + float base = saturated_f * scale_to_zero; + + const uint32_t w = std::bit_cast(f); + const uint32_t shl1_w = w + w; + const uint32_t sign = w & UINT32_C(0x80000000); + uint32_t bias = shl1_w & UINT32_C(0xFF000000); + if (bias < UINT32_C(0x71000000)) { + bias = UINT32_C(0x71000000); + } + + base = std::bit_cast((bias >> 1) + UINT32_C(0x07800000)) + base; + const uint32_t bits = std::bit_cast(base); + const uint32_t exp_bits = (bits >> 13) & UINT32_C(0x00007C00); + const uint32_t mantissa_bits = bits & UINT32_C(0x00000FFF); + const uint32_t nonsign = exp_bits + mantissa_bits; + return (sign >> 16) | (shl1_w > UINT32_C(0xFF000000) ? UINT16_C(0x7E00) : nonsign); +} + +static constexpr float fp16_ieee_to_fp32_value(uint16_t h) { + /* + * Extend the half-precision floating-point number to 32 bits and shift to + * the upper part of the 32-bit word: + * +---+-----+------------+-------------------+ + * | S |EEEEE|MM MMMM MMMM|0000 0000 0000 0000| + * +---+-----+------------+-------------------+ + * Bits 31 26-30 16-25 0-15 + * + * S - sign bit, E - bits of the biased exponent, M - bits of the mantissa, + * 0 - zero bits. + */ + const uint32_t w = (uint32_t)h << 16; + /* + * Extract the sign of the input number into the high bit of the 32-bit + * word: + * + * +---+----------------------------------+ + * | S |0000000 00000000 00000000 00000000| + * +---+----------------------------------+ + * Bits 31 0-31 + */ + const uint32_t sign = w & UINT32_C(0x80000000); + /* + * Extract mantissa and biased exponent of the input number into the high + * bits of the 32-bit word: + * + * +-----+------------+---------------------+ + * |EEEEE|MM MMMM MMMM|0 0000 0000 0000 0000| + * +-----+------------+---------------------+ + * Bits 27-31 17-26 0-16 + */ + const uint32_t two_w = w + w; + + /* + * Shift mantissa and exponent into bits 23-28 and bits 13-22 so they become + * mantissa and exponent of a single-precision floating-point number: + * + * S|Exponent | Mantissa + * +-+---+-----+------------+----------------+ + * |0|000|EEEEE|MM MMMM MMMM|0 0000 0000 0000| + * +-+---+-----+------------+----------------+ + * Bits | 23-31 | 0-22 + * + * Next, there are some adjustments to the exponent: + * - The exponent needs to be corrected by the difference in exponent bias + * between single-precision and half-precision + * formats (0x7F - 0xF = 0x70) + * - Inf and NaN values in the inputs should become Inf and NaN values after + * conversion to the single-precision number. + * Therefore, if the biased exponent of the half-precision input was 0x1F + * (max possible value), the biased exponent + * of the single-precision output must be 0xFF (max possible value). We do + * this correction in two steps: + * - First, we adjust the exponent by (0xFF - 0x1F) = 0xE0 (see exp_offset + * below) rather than by 0x70 suggested + * by the difference in the exponent bias (see above). + * - Then we multiply the single-precision result of exponent adjustment + * by 2**(-112) to reverse the effect of + * exponent adjustment by 0xE0 less the necessary exponent adjustment by + * 0x70 due to difference in exponent bias. + * The floating-point multiplication hardware would ensure than Inf and + * NaN would retain their value on at least + * partially IEEE754-compliant implementations. + * + * Note that the above operations do not handle denormal inputs (where + * biased exponent == 0). However, they also do not operate on denormal + * inputs, and do not produce denormal results. + */ + const uint32_t exp_offset = UINT32_C(0xE0) << 23; + const float exp_scale = std::bit_cast(UINT32_C(0x7800000)); + const float normalized_value = std::bit_cast((two_w >> 4) + exp_offset) * exp_scale; + + /* + * Convert denormalized half-precision inputs into single-precision results + * (always normalized). + * Zero inputs are also handled here. + * + * In a denormalized number the biased exponent is zero, and mantissa has + * on-zero bits. + * First, we shift mantissa into bits 0-9 of the 32-bit word. + * + * zeros | mantissa + * +---------------------------+------------+ + * |0000 0000 0000 0000 0000 00|MM MMMM MMMM| + * +---------------------------+------------+ + * Bits 10-31 0-9 + * + * Now, remember that denormalized half-precision numbers are represented + * as: + * FP16 = mantissa * 2**(-24). + * The trick is to construct a normalized single-precision number with the + * same mantissa and thehalf-precision input + * and with an exponent which would scale the corresponding mantissa bits + * to 2**(-24). + * A normalized single-precision floating-point number is represented as: + * FP32 = (1 + mantissa * 2**(-23)) * 2**(exponent - 127) + * Therefore, when the biased exponent is 126, a unit change in the mantissa + * of the input denormalized half-precision + * number causes a change of the constructud single-precision number by + * 2**(-24), i.e. the same ammount. + * + * The last step is to adjust the bias of the constructed single-precision + * number. When the input half-precision number + * is zero, the constructed single-precision number has the value of + * FP32 = 1 * 2**(126 - 127) = 2**(-1) = 0.5 + * Therefore, we need to subtract 0.5 from the constructed single-precision + * number to get the numerical equivalent of + * the input half-precision number. + */ + const uint32_t magic_mask = UINT32_C(126) << 23; + const float magic_bias = 0.5f; + const float denormalized_value = std::bit_cast((two_w >> 17) | magic_mask) - magic_bias; + + /* + * - Choose either results of conversion of input as a normalized number, or + * as a denormalized number, depending on the + * input exponent. The variable two_w contains input exponent in bits + * 27-31, therefore if its smaller than 2**27, the + * input is either a denormal number, or zero. + * - Combine the result of conversion of exponent and mantissa with the sign + * of the input number. + */ + const uint32_t denormalized_cutoff = UINT32_C(1) << 27; + const uint32_t result = + sign | (two_w < denormalized_cutoff ? std::bit_cast(denormalized_value) + : std::bit_cast(normalized_value)); + return std::bit_cast(result); +#endif +} diff --git a/include/flashinfer/frag_layout_swizzle.cuh b/include/flashinfer/frag_layout_swizzle.cuh new file mode 100644 index 0000000000000000000000000000000000000000..39cf92bcd96bf1cd015b7db7382fb73c61f03110 --- /dev/null +++ b/include/flashinfer/frag_layout_swizzle.cuh @@ -0,0 +1,41 @@ +/* + * Copyright (c) 2024 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef FLASHINFER_FRAG_LAYOUT_SWIZZLE_CUH_ +#define FLASHINFER_FRAG_LAYOUT_SWIZZLE_CUH_ + +#include + +#include + +__device__ __forceinline__ uint32_t frag_layout_swizzle_16b_to_8b(uint32_t x) { + uint32_t tmp = __shfl_xor_sync(0xffffffff, x, 0x1); + x = __byte_perm(x, tmp, ((threadIdx.x & 0x1) == 0) ? 0x5410 : 0x3276); + tmp = __shfl_xor_sync(0xffffffff, x, 0x2); + x = __byte_perm(x, tmp, ((threadIdx.x & 0x2) == 0) ? 0x5410 : 0x3276); + return x; +} + +__device__ __forceinline__ uint32_t frag_layout_swizzle_16b_to_8b_trans(uint32_t x) { + uint32_t tmp = __shfl_xor_sync(0xffffffff, x, 0x4); + x = __byte_perm(x, tmp, ((threadIdx.x & 0x4) == 0) ? 0x6420 : 0x3175); + tmp = __shfl_xor_sync(0xffffffff, x, 0x8); + x = __byte_perm(x, tmp, ((threadIdx.x & 0x8) == 0) ? 0x5410 : 0x3276); + tmp = __shfl_xor_sync(0xffffffff, x, 0x10); + x = __byte_perm(x, tmp, ((threadIdx.x & 0x10) == 0) ? 0x5410 : 0x3276); + return x; +} + +#endif // FLASHINFER_FRAG_LAYOUT_SWIZZLE_CUH_ diff --git a/include/flashinfer/gemm/bmm_fp8.cuh b/include/flashinfer/gemm/bmm_fp8.cuh new file mode 100644 index 0000000000000000000000000000000000000000..7778934160c5da191498ffde08ef0c27a872a274 --- /dev/null +++ b/include/flashinfer/gemm/bmm_fp8.cuh @@ -0,0 +1,203 @@ +/* + * Copyright (c) 2024 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef FLASHINFER_GEMM_BMM_FP8_CUH_ +#define FLASHINFER_GEMM_BMM_FP8_CUH_ + +#include +#include + +#include +#include +#include + +#include "../exception.h" + +#define FLASHINFER_CUBLAS_CHECK(EXPR) \ + { \ + cublasStatus_t e = (EXPR); \ + FLASHINFER_CHECK(e == CUBLAS_STATUS_SUCCESS, \ + "CUBLAS Error: " + std::string(cublasGetStatusString(e))); \ + } + +#ifndef NDEBUG +#define FLASHINFER_CUBLAS_CALL(EXPR) \ + { \ + cublasStatus_t e = (EXPR); \ + if (e != CUBLAS_STATUS_SUCCESS) { \ + std::cerr << "CUBLAS Error: " << cublasGetStatusString(e) << " (" << e << ") " << __FILE__ \ + << ": line " << __LINE__ << " at function " << #EXPR << std::endl; \ + return e; \ + } \ + } +#else +#define FLASHINFER_CUBLAS_CALL(EXPR) \ + { \ + cublasStatus_t e = (EXPR); \ + if (e != CUBLAS_STATUS_SUCCESS) { \ + return e; \ + } \ + } +#endif + +namespace flashinfer { + +namespace bmm_fp8 { + +template +struct CuBlasLtDeleter { + void operator()(T* x) { + if (x != nullptr) { + FLASHINFER_CUBLAS_CHECK(destructor(x)); + } + } +}; + +template +class CuBlasLtDescriptor { + public: + T* descriptor() const { return descriptor_.get(); } + T* descriptor() { return descriptor_.get(); } + + protected: + std::unique_ptr> descriptor_; +}; + +class CuBlasLtMatmulDescriptor + : public CuBlasLtDescriptor { + public: + CuBlasLtMatmulDescriptor(cublasComputeType_t compute_type, cudaDataType_t scale_type) { + cublasLtMatmulDesc_t raw_descriptor = nullptr; + FLASHINFER_CUBLAS_CHECK(cublasLtMatmulDescCreate(&raw_descriptor, compute_type, scale_type)); + descriptor_.reset(raw_descriptor); + } + template + inline void setAttribute(cublasLtMatmulDescAttributes_t attr, const T value) { + FLASHINFER_CUBLAS_CHECK( + ::cublasLtMatmulDescSetAttribute(descriptor(), attr, &value, sizeof(T))); + } +}; + +class CuBlasLtMatrixLayout + : public CuBlasLtDescriptor { + public: + CuBlasLtMatrixLayout(cudaDataType_t type, uint64_t rows, uint64_t cols, int64_t ld, + bool t = false) { + cublasLtMatrixLayout_t raw_descriptor = nullptr; + FLASHINFER_CUBLAS_CHECK( + cublasLtMatrixLayoutCreate(&raw_descriptor, type, t ? cols : rows, t ? rows : cols, ld)); + descriptor_.reset(raw_descriptor); + } + template + inline void setAttribute(cublasLtMatrixLayoutAttribute_t attr, const T value) { + FLASHINFER_CUBLAS_CHECK( + ::cublasLtMatrixLayoutSetAttribute(descriptor(), attr, &value, sizeof(T))); + } +}; + +class CuBlasLtMatmulPreference : public CuBlasLtDescriptor { + public: + CuBlasLtMatmulPreference() { + cublasLtMatmulPreference_t raw_descriptor = nullptr; + FLASHINFER_CUBLAS_CHECK(cublasLtMatmulPreferenceCreate(&raw_descriptor)); + descriptor_.reset(raw_descriptor); + } + template + inline void setAttribute(cublasLtMatmulPreferenceAttributes_t attr, const T value) { + FLASHINFER_CUBLAS_CHECK( + ::cublasLtMatmulPreferenceSetAttribute(descriptor(), attr, &value, sizeof(T))); + } +}; + +template +cudaDataType_t get_cuda_data_type() { + if constexpr (std::is_same_v) { + return CUDA_R_8F_E4M3; + } else if constexpr (std::is_same_v) { + return CUDA_R_8F_E5M2; + } else if constexpr (std::is_same_v) { + return CUDA_R_16BF; + } else if constexpr (std::is_same_v) { + return CUDA_R_16F; + } else { + FLASHINFER_ERROR("Unsupported type"); + } +} + +template +cublasStatus_t bmm_fp8_internal_cublaslt(void* workspace, size_t workspace_size_in_bytes, + const AT* A, const BT* B, DT* D, int batch_size, int m, + int n, int k, const float* A_scale, const float* B_scale, + cublasLtHandle_t lt_handle, cudaStream_t stream) { + const void* A_scale_ptr = static_cast(A_scale); + const void* B_scale_ptr = static_cast(B_scale); + auto matmul_desp = CuBlasLtMatmulDescriptor(CUBLAS_COMPUTE_32F, CUDA_R_32F); + matmul_desp.setAttribute(CUBLASLT_MATMUL_DESC_TRANSA, CUBLAS_OP_T); + matmul_desp.setAttribute(CUBLASLT_MATMUL_DESC_TRANSB, CUBLAS_OP_N); + int8_t fast_accum = 1; + matmul_desp.setAttribute(CUBLASLT_MATMUL_DESC_FAST_ACCUM, fast_accum); + + matmul_desp.setAttribute(CUBLASLT_MATMUL_DESC_A_SCALE_POINTER, A_scale_ptr); + matmul_desp.setAttribute(CUBLASLT_MATMUL_DESC_B_SCALE_POINTER, B_scale_ptr); + + cudaDataType_t a_type = get_cuda_data_type(); + cudaDataType_t b_type = get_cuda_data_type(); + cudaDataType_t d_type = get_cuda_data_type

(); + if (std::is_same_v && std::is_same_v) { + FLASHINFER_ERROR("Unsupported combination: both A and B are e5m2"); + } + + auto a_desp = CuBlasLtMatrixLayout(a_type, m, k, k, true); + auto b_desp = CuBlasLtMatrixLayout(b_type, k, n, k); + auto d_desp = CuBlasLtMatrixLayout(d_type, m, n, m); + + if (batch_size > 1) { + int64_t stride_a = m * k; + int64_t stride_b = k * n; + int64_t stride_d = m * n; + a_desp.setAttribute(CUBLASLT_MATRIX_LAYOUT_BATCH_COUNT, batch_size); + a_desp.setAttribute(CUBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET, stride_a); + b_desp.setAttribute(CUBLASLT_MATRIX_LAYOUT_BATCH_COUNT, batch_size); + b_desp.setAttribute(CUBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET, stride_b); + d_desp.setAttribute(CUBLASLT_MATRIX_LAYOUT_BATCH_COUNT, batch_size); + d_desp.setAttribute(CUBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET, stride_d); + } + + CuBlasLtMatmulPreference preference; + preference.setAttribute(CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, workspace_size_in_bytes); + cublasLtMatmulHeuristicResult_t heuristic_result = {}; + int returned_result = 0; + FLASHINFER_CUBLAS_CALL(cublasLtMatmulAlgoGetHeuristic( + lt_handle, matmul_desp.descriptor(), a_desp.descriptor(), b_desp.descriptor(), + d_desp.descriptor(), d_desp.descriptor(), preference.descriptor(), 1, &heuristic_result, + &returned_result)); + if (returned_result == 0) { + return CUBLAS_STATUS_NOT_SUPPORTED; + } + + const float alpha = 1.0f; + const float beta = 0.0f; + FLASHINFER_CUBLAS_CALL(cublasLtMatmul( + lt_handle, matmul_desp.descriptor(), &alpha, A, a_desp.descriptor(), B, b_desp.descriptor(), + &beta, nullptr, d_desp.descriptor(), D, d_desp.descriptor(), &heuristic_result.algo, + workspace, workspace_size_in_bytes, stream)); + return CUBLAS_STATUS_SUCCESS; +} + +} // namespace bmm_fp8 +} // namespace flashinfer + +#endif // FLASHINFER_GEMM_BMM_FP8_CUH_ diff --git a/include/flashinfer/gemm/cutlass_gemm_configs.h b/include/flashinfer/gemm/cutlass_gemm_configs.h new file mode 100644 index 0000000000000000000000000000000000000000..993583ddefcdaf01dd0008a80aa0ebb4afa52768 --- /dev/null +++ b/include/flashinfer/gemm/cutlass_gemm_configs.h @@ -0,0 +1,425 @@ +/* + * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef FLASHINFER_CUTLASS_GEMM_CONFIG_H_ +#define FLASHINFER_CUTLASS_GEMM_CONFIG_H_ + +#include +#include +#include +#include + +#include "cute/tensor.hpp" + +namespace flashinfer { +namespace gemm { + +// Note: The shapes are in the format MxNxK. The K shape of the runtime config MUST match the K +// shape +// in the kernel layout details when doing weight only quantization. +enum class CutlassTileConfig { + // Signals that we should run heuristics do choose a config + Undefined, + + // Signals that we should run heuristics do choose a config + ChooseWithHeuristic, + + // SiMT config + CtaShape128x128x8_WarpShape64x64x8, + + // TensorCore configs CTA_N = 128, CTA_K = 64 + // Warp configs for M=16 + CtaShape16x128x64_WarpShape16x32x64, + // Warp configs for M=32 + CtaShape32x128x64_WarpShape32x32x64, + + // Warp configs for M=64 + CtaShape64x128x64_WarpShape32x64x64, + CtaShape64x64x128_WarpShape32x64x64, + CtaShape64x128x64_WarpShape64x32x64, + + // Warp configs for M=128 + CtaShape128x64x64_WarpShape64x32x64, + CtaShape128x128x64_WarpShape64x32x64, + CtaShape128x128x64_WarpShape64x64x64, + CtaShape128x128x64_WarpShape128x32x64, + CtaShape128x256x64_WarpShape64x64x64, + + // Warp configs for M=256 + CtaShape256x128x64_WarpShape64x64x64, + + // TensorCore config CTA_N = 64, CTA_K = 128 + CtaShape128x64x128_WarpShape64x32x128, + + // TensorCore config CTA_N = 256, CTA_K = 64 + CtaShape16x256x64_WarpShape16x64x64, + + // TensorCore config CTA_N = 256, CTA_K = 128 + CtaShape16x256x128_WarpShape16x64x128 + +}; + +enum class SplitKStyle { + NO_SPLIT_K, + SPLIT_K_SERIAL, + STREAM_K, // Sm80+ + // SPLIT_K_PARALLEL // Not supported yet +}; + +enum class CutlassTileConfigSM90 { + // Signals that we should run heuristics do choose a config + Undefined, + + // Signals that we should run heuristics do choose a config + ChooseWithHeuristic, + + // CTA configs for M=64 + CtaShape64x16x128B, + CtaShape64x32x128B, + CtaShape64x64x128B, + CtaShape64x128x128B, + CtaShape64x256x128B, + + // CTA configs for M=128 + CtaShape128x16x128B, + CtaShape128x32x128B, + CtaShape128x64x128B, + CtaShape128x128x128B, + CtaShape128x256x128B, + + // CTA configs for M=128 + CtaShape256x128x128B, +}; + +enum class CutlassTileConfigSM100 { + // Signals that we should run heuristics do choose a config + Undefined, + + // Signals that we should run heuristics do choose a config + ChooseWithHeuristic, + + /* + * Grouped GEMM + */ + // M=64 + CtaShape64x32x128B, + CtaShape64x64x128B, + CtaShape64x128x128B, + CtaShape64x256x128B, + + // M=128 + CtaShape128x8x256B, + CtaShape128x16x128B, + CtaShape128x32x128B, + CtaShape128x64x128B, + CtaShape128x128x128B, + CtaShape128x256x128B, + CtaShape128x128x256B, + CtaShape128x256x256B, + + // M=256 + CtaShape256x64x128B, + CtaShape256x128x128B, + CtaShape256x256x128B, +}; + +enum class CutlassTileConfigSM120 { + // Signals that we should run heuristics do choose a config + Undefined, + + // Signals that we should run heuristics do choose a config + ChooseWithHeuristic, + + CtaShape128x128x128B, + CtaShape128x128x64B, + CtaShape256x128x64B, + CtaShape128x256x64B, + CtaShape128x128x256B, + CtaShape256x128x128B, +}; + +enum class MainloopScheduleType { + AUTO, // Automatically selects between pingpong and cooperative schedules on Hopper. On older + // architectures, this defaults to the "legacy" main loop schedule. + PINGPONG, + COOPERATIVE, + WARPSPECIALIZED +}; + +static auto get_mainloop_schedule_name(MainloopScheduleType schedule) { + if (schedule == MainloopScheduleType::AUTO) { + return "auto"; + } else if (schedule == MainloopScheduleType::PINGPONG) { + return "pingpong"; + } else if (schedule == MainloopScheduleType::COOPERATIVE) { + return "cooperative"; + } else if (schedule == MainloopScheduleType::WARPSPECIALIZED) { + return "warpspecialized"; + } + return "unknown schedule"; +} + +enum class EpilogueScheduleType { + AUTO, // Automatically chooses an epilogue schedule compatible with the selected main loop + // schedule for Hopper. For architectures older than hopper, the epilogue is always + // performed by the same thread block as the main loop. +}; + +enum class TileShape { + TileShape_64x16x128, + TileShape_64x32x128, + TileShape_64x64x128, + TileShape_64x128x128, + TileShape_64x256x128, + TileShape_64x512x128, + TileShape_128x16x128, + TileShape_128x32x128, + TileShape_128x64x128, + TileShape_128x128x128, + TileShape_128x256x128 +}; + +template +constexpr auto get_tile_shape() { + using namespace cute; + if constexpr (Shape_MNK == TileShape::TileShape_64x16x128) { + return cute::Shape<_64, _16, _128>{}; + } else if constexpr (Shape_MNK == TileShape::TileShape_64x32x128) { + return cute::Shape<_64, _32, _128>{}; + } else if constexpr (Shape_MNK == TileShape::TileShape_64x64x128) { + return cute::Shape<_64, _64, _128>{}; + } else if constexpr (Shape_MNK == TileShape::TileShape_64x128x128) { + return cute::Shape<_64, _128, _128>{}; + } else if constexpr (Shape_MNK == TileShape::TileShape_64x256x128) { + return cute::Shape<_64, _256, _128>{}; + } else if constexpr (Shape_MNK == TileShape::TileShape_64x512x128) { + return cute::Shape<_64, _512, _128>{}; + } else if constexpr (Shape_MNK == TileShape::TileShape_128x16x128) { + return cute::Shape<_128, _16, _128>{}; + } else if constexpr (Shape_MNK == TileShape::TileShape_128x32x128) { + return cute::Shape<_128, _32, _128>{}; + } else if constexpr (Shape_MNK == TileShape::TileShape_128x64x128) { + return cute::Shape<_128, _64, _128>{}; + } else if constexpr (Shape_MNK == TileShape::TileShape_128x128x128) { + return cute::Shape<_128, _128, _128>{}; + } else if constexpr (Shape_MNK == TileShape::TileShape_128x256x128) { + return cute::Shape<_128, _256, _128>{}; + } +} + +static auto get_tile_shape_name(TileShape Shape_MNK) { + if (Shape_MNK == TileShape::TileShape_64x16x128) { + return "64x16x128"; + } else if (Shape_MNK == TileShape::TileShape_64x32x128) { + return "64x32x128"; + } else if (Shape_MNK == TileShape::TileShape_64x64x128) { + return "64x64x128"; + } else if (Shape_MNK == TileShape::TileShape_64x128x128) { + return "64x128x128"; + } else if (Shape_MNK == TileShape::TileShape_64x256x128) { + return "64x256x128"; + } else if (Shape_MNK == TileShape::TileShape_64x512x128) { + return "64x512x128"; + } else if (Shape_MNK == TileShape::TileShape_128x16x128) { + return "128x16x128"; + } else if (Shape_MNK == TileShape::TileShape_128x32x128) { + return "128x32x128"; + } else if (Shape_MNK == TileShape::TileShape_128x64x128) { + return "128x64x128"; + } else if (Shape_MNK == TileShape::TileShape_128x128x128) { + return "128x128x128"; + } else if (Shape_MNK == TileShape::TileShape_128x256x128) { + return "128x256x128"; + } + return "Unknown shape"; +} + +enum class ClusterShape { + ClusterShape_1x1x1, + ClusterShape_2x1x1, + ClusterShape_1x2x1, + ClusterShape_2x2x1, + ClusterShape_1x4x1, + ClusterShape_4x2x1, + ClusterShape_2x4x1, + ClusterShape_4x4x1, + ClusterShape_1x8x1, + ClusterShape_8x1x1 +}; + +static auto get_cluster_shape_name(ClusterShape Shape_MNK) { + if (Shape_MNK == ClusterShape::ClusterShape_1x1x1) { + return "1x1x1"; + } else if (Shape_MNK == ClusterShape::ClusterShape_2x1x1) { + return "2x1x1"; + } else if (Shape_MNK == ClusterShape::ClusterShape_1x2x1) { + return "1x2x1"; + } else if (Shape_MNK == ClusterShape::ClusterShape_2x2x1) { + return "2x2x1"; + } else if (Shape_MNK == ClusterShape::ClusterShape_1x8x1) { + return "1x8x1"; + } else if (Shape_MNK == ClusterShape::ClusterShape_8x1x1) { + return "8x1x1"; + } + return "Unknown shape"; +} + +template +constexpr auto get_cluster_shape() { + using namespace cute; + if constexpr (Shape_MNK == ClusterShape::ClusterShape_1x1x1) { + return cute::Shape<_1, _1, _1>{}; + } else if constexpr (Shape_MNK == ClusterShape::ClusterShape_2x1x1) { + return cute::Shape<_2, _1, _1>{}; + } else if constexpr (Shape_MNK == ClusterShape::ClusterShape_1x2x1) { + return cute::Shape<_1, _2, _1>{}; + } else if constexpr (Shape_MNK == ClusterShape::ClusterShape_2x2x1) { + return cute::Shape<_2, _2, _1>{}; + } else if constexpr (Shape_MNK == ClusterShape::ClusterShape_1x8x1) { + return cute::Shape<_1, _8, _1>{}; + } else if constexpr (Shape_MNK == ClusterShape::ClusterShape_8x1x1) { + return cute::Shape<_8, _1, _1>{}; + } +} + +struct CutlassGemmConfig { + enum CandidateConfigTypeParam : int { + NONE = 0, + WEIGHT_ONLY = 1u << 0, + SIMT_ONLY = 1u << 1, + INT8_ONLY = 1u << 2, + HOPPER = 1u << 3, + BLACKWELL = 1u << 4, + GROUPED_GEMM = 1u << 5, + FP8_ONLY = 1u << 6, + FP4_ONLY = 1u << 7 + }; + + CutlassTileConfig tile_config_sm80 = CutlassTileConfig::ChooseWithHeuristic; + SplitKStyle split_k_style = SplitKStyle::NO_SPLIT_K; + int split_k_factor = -1; + int stages = -1; + + // config options for sm90 + CutlassTileConfigSM90 tile_config_sm90 = CutlassTileConfigSM90::ChooseWithHeuristic; + CutlassTileConfigSM100 tile_config_sm100 = CutlassTileConfigSM100::ChooseWithHeuristic; + CutlassTileConfigSM120 tile_config_sm120 = CutlassTileConfigSM120::ChooseWithHeuristic; + MainloopScheduleType mainloop_schedule = MainloopScheduleType::AUTO; + EpilogueScheduleType epilogue_schedule = EpilogueScheduleType::AUTO; + ClusterShape cluster_shape = ClusterShape::ClusterShape_1x1x1; + bool enableCudaKernel = false; + int sm_version = 80; // Use 80 as a catch all for <90 + bool is_tma_warp_specialized = false; + + CutlassGemmConfig() = default; + + CutlassGemmConfig(CutlassTileConfig tile_config, SplitKStyle split_k_style, int split_k_factor, + int stages) + : tile_config_sm80(tile_config), + split_k_style(split_k_style), + split_k_factor(split_k_factor), + stages(stages), + sm_version(80) {} + + CutlassGemmConfig(CutlassTileConfigSM90 tile_config_sm90, MainloopScheduleType mainloop_schedule, + EpilogueScheduleType epilogue_schedule, ClusterShape cluster_shape) + : tile_config_sm90(tile_config_sm90), + mainloop_schedule(mainloop_schedule), + epilogue_schedule(epilogue_schedule), + cluster_shape(cluster_shape), + sm_version(90), + is_tma_warp_specialized(true) {} + + CutlassGemmConfig(CutlassTileConfigSM100 tile_config_sm100, + MainloopScheduleType mainloop_schedule, EpilogueScheduleType epilogue_schedule, + ClusterShape cluster_shape) + : tile_config_sm100(tile_config_sm100), + mainloop_schedule(mainloop_schedule), + epilogue_schedule(epilogue_schedule), + cluster_shape(cluster_shape), + sm_version(100), + is_tma_warp_specialized(true) {} + + CutlassGemmConfig(CutlassTileConfigSM120 tile_config_sm120, + MainloopScheduleType mainloop_schedule, EpilogueScheduleType epilogue_schedule, + ClusterShape cluster_shape) + : tile_config_sm120(tile_config_sm120), + mainloop_schedule(mainloop_schedule), + epilogue_schedule(epilogue_schedule), + cluster_shape(cluster_shape), + sm_version(120), + is_tma_warp_specialized(true) {} + + int getTileConfigAsInt() const { + if (sm_version == 120) return (int)tile_config_sm120; + if (sm_version >= 100) return (int)tile_config_sm100; + if (sm_version == 90) return (int)tile_config_sm90; + if (sm_version < 90) return (int)tile_config_sm80; + assert(false && "Invalid SM version"); + return -1; + } + + std::string toString() const { + std::stringstream tactic; + tactic << "Cutlass GEMM Tactic"; + if (is_tma_warp_specialized) { + assert(sm_version >= 90 && "Invalid cutlass GEMM config"); + tactic << "\n\tstyle=TMA Warp Specialized" + << "\n\tsm: " << sm_version << "\n\ttile shape ID: " << getTileConfigAsInt() + << "\n\tcluster shape ID: " << (int)cluster_shape + << "\n\tmainloop sched: " << (int)mainloop_schedule + << "\n\tepi sched: " << (int)epilogue_schedule + << "\n\tenable cuda kernel: " << (enableCudaKernel ? "true" : "false"); + } else if (tile_config_sm80 != flashinfer::gemm::CutlassTileConfig::ChooseWithHeuristic) { + assert(sm_version < 90 && "Invalid cutlass GEMM config"); + tactic << "\n\tstyle=compatible" + << "\n\ttile shape ID: " << (int)tile_config_sm80 << "\n\tstages: " << (int)stages + << "\n\tsplit k: " << (int)split_k_factor + << "\n\tenable cuda kernel: " << (enableCudaKernel ? "true" : "false"); + } else if (enableCudaKernel) { + tactic << "\n\tenable cuda kernel: " << (enableCudaKernel ? "true" : "false"); + } else { + tactic << "\n\tundefined"; + } + tactic << "\n"; + return tactic.str(); + } +}; + +inline std::ostream& operator<<(std::ostream& out, CutlassGemmConfig const& config) { + // clang-format off + if (config.is_tma_warp_specialized) + { + out << "tile_config_sm90_enum: " << config.getTileConfigAsInt() + << ", mainloop_schedule_enum: " << int(config.mainloop_schedule) + << ", epilogue_schedule_enum: " << int(config.epilogue_schedule) + << ", cluster_shape_enum: " << int(config.cluster_shape) + << ", enable_cuda_kernel: " << (config.enableCudaKernel ? "true" : "false"); + } + else + { + out << "tile_config_enum: " << config.getTileConfigAsInt() + << ", split_k_style_enum: " << int(config.split_k_style) + << ", split_k_factor: " << config.split_k_factor + << ", stages: " << config.stages + << ", enable_cuda_kernel: " << (config.enableCudaKernel ? "true" : "false"); + } + // clang-format on + return out; +} + +} // namespace gemm +} // namespace flashinfer +#endif // FLASHINFER_CUTLASS_GEMM_CONFIG_H_ diff --git a/include/flashinfer/gemm/fp4_gemm_cutlass.h b/include/flashinfer/gemm/fp4_gemm_cutlass.h new file mode 100644 index 0000000000000000000000000000000000000000..d4f21b84c8a5d0711f9169db4ed783ea4a971acb --- /dev/null +++ b/include/flashinfer/gemm/fp4_gemm_cutlass.h @@ -0,0 +1,91 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights + * reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#ifndef FLASHINFER_FP4_GEMM_CUTLASS_H_ +#define FLASHINFER_FP4_GEMM_CUTLASS_H_ + +#include + +#include + +#include "flashinfer/gemm/cutlass_gemm_configs.h" + +namespace flashinfer { +namespace gemm { + +/* + This runner supports: + FP4 inputs (A and B) + float blockwise scaling factor + float alpha scalings + T output (D) where T = {float, half, __nv_bfloat16} + + Activations, biases and outputs are all assumed to be row-major. + Weights are assumed to be column-major. + Block scaling factor are interleaved. +*/ + +class CutlassFp4GemmRunnerInterface { + public: + CutlassFp4GemmRunnerInterface() {} + + virtual ~CutlassFp4GemmRunnerInterface() {} + + virtual void gemm(void* D, void const* A, void const* B, void const* input_sf, + void const* weight_sf, float const* global_sf, int m, int n, int k, + int batch_count, CutlassGemmConfig gemmConfig, char* workspace, + const size_t workspaceBytes, cudaStream_t stream) = 0; + + // Returns desired workspace size in bytes. + virtual size_t getWorkspaceSize(int const m, int const n, int const k, int batch_count) = 0; + + virtual std::vector getConfigs() const = 0; +}; + +enum class FP4GemmType { + W4A4_NVFP4_NVFP4, +}; + +template +class CutlassFp4GemmRunner : public virtual CutlassFp4GemmRunnerInterface { + public: + CutlassFp4GemmRunner(); + ~CutlassFp4GemmRunner(); + + void gemm(void* D, void const* A, void const* B, void const* input_sf, void const* weight_sf, + float const* global_sf, int m, int n, int k, int batch_count, + CutlassGemmConfig gemmConfig, char* workspace, const size_t workspaceBytes, + cudaStream_t stream) override; + + // Returns desired workspace size in bytes. + size_t getWorkspaceSize(int const m, int const n, int const k, int const batch_count) override; + + std::vector getConfigs() const override; + + private: + size_t dispatchToArch(T* D, void const* A, void const* B, void const* input_sf, + void const* weight_sf, float const* global_sf, int m, int n, int k, + int batch_count, CutlassGemmConfig gemmConfig, char* workspace, + const size_t workspaceBytes, cudaStream_t stream, int* occupancy = nullptr); + + size_t getWorkspaceSizeImpl(int const m, int const n, int const k, int const batch_count); +}; + +} // namespace gemm +} // namespace flashinfer + +#endif // FLASHINFER_FP4_GEMM_CUTLASS_H_ diff --git a/include/flashinfer/gemm/fp4_gemm_cutlass_template.h b/include/flashinfer/gemm/fp4_gemm_cutlass_template.h new file mode 100644 index 0000000000000000000000000000000000000000..979a660eadf97b9ede953f3c77c4d05bc82e54b4 --- /dev/null +++ b/include/flashinfer/gemm/fp4_gemm_cutlass_template.h @@ -0,0 +1,284 @@ +/* + * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#ifndef FLASHINFER_FP4_GEMM_CUTLASS_TEMPLATE_H_ +#define FLASHINFER_FP4_GEMM_CUTLASS_TEMPLATE_H_ + +#ifndef _WIN32 +#pragma GCC diagnostic push +#pragma GCC diagnostic ignored "-Wstrict-aliasing" +#endif // #ifndef _WIN32 + +#include "cutlass/arch/arch.h" +#include "cutlass/cutlass.h" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/gemm.h" +#include "flashinfer/gemm/cutlass_gemm_configs.h" + +#ifndef _WIN32 +#pragma GCC diagnostic pop +#endif // #ifndef _WIN32 + +#include "flashinfer/gemm/fp4_gemm_cutlass.h" +#include "fp4_gemm_template_sm100.h" + +namespace flashinfer { +namespace gemm { +using namespace cute; + +template +size_t dispatchNVFP4xNVFP4GemmClusterShapeSm100(T* D, void const* A, void const* B, + void const* input_sf, void const* weight_sf, + float const* global_sf, int m, int n, int k, + int batch_count, CutlassGemmConfig gemmConfig, + char* workspace, const size_t workspaceBytes, + cudaStream_t stream, int* occupancy = nullptr) { + switch (gemmConfig.cluster_shape) { + case ClusterShape::ClusterShape_1x1x1: + return genericFp4GemmKernelLauncher, cute::Int<1>, + cute::Int<1>, _1SM>( + D, A, B, input_sf, weight_sf, global_sf, m, n, k, batch_count, gemmConfig, workspace, + workspaceBytes, stream, occupancy); + break; + case ClusterShape::ClusterShape_2x1x1: + return genericFp4GemmKernelLauncher, cute::Int<1>, + cute::Int<1>, _2SM>( + D, A, B, input_sf, weight_sf, global_sf, m, n, k, batch_count, gemmConfig, workspace, + workspaceBytes, stream, occupancy); + break; + case ClusterShape::ClusterShape_1x2x1: + return genericFp4GemmKernelLauncher, cute::Int<2>, + cute::Int<1>, _1SM>( + D, A, B, input_sf, weight_sf, global_sf, m, n, k, batch_count, gemmConfig, workspace, + workspaceBytes, stream, occupancy); + break; + case ClusterShape::ClusterShape_2x2x1: + return genericFp4GemmKernelLauncher, cute::Int<2>, + cute::Int<1>, _2SM>( + D, A, B, input_sf, weight_sf, global_sf, m, n, k, batch_count, gemmConfig, workspace, + workspaceBytes, stream, occupancy); + break; + case ClusterShape::ClusterShape_1x4x1: + return genericFp4GemmKernelLauncher, cute::Int<4>, + cute::Int<1>, _1SM>( + D, A, B, input_sf, weight_sf, global_sf, m, n, k, batch_count, gemmConfig, workspace, + workspaceBytes, stream, occupancy); + break; + case ClusterShape::ClusterShape_4x2x1: + return genericFp4GemmKernelLauncher, cute::Int<2>, + cute::Int<1>, _2SM>( + D, A, B, input_sf, weight_sf, global_sf, m, n, k, batch_count, gemmConfig, workspace, + workspaceBytes, stream, occupancy); + break; + case ClusterShape::ClusterShape_2x4x1: + return genericFp4GemmKernelLauncher, cute::Int<4>, + cute::Int<1>, _2SM>( + D, A, B, input_sf, weight_sf, global_sf, m, n, k, batch_count, gemmConfig, workspace, + workspaceBytes, stream, occupancy); + break; + case ClusterShape::ClusterShape_4x4x1: + return genericFp4GemmKernelLauncher, cute::Int<4>, + cute::Int<1>, _2SM>( + D, A, B, input_sf, weight_sf, global_sf, m, n, k, batch_count, gemmConfig, workspace, + workspaceBytes, stream, occupancy); + break; + default: + throw std::runtime_error( + "[Error][FP4][dispatch_gemm_cluster_shape] Config is invalid for FP4 GEMM."); + break; + } +} + +template +size_t dispatchNVFP4xNVFP4GemmCTAShapeSm100(T* D, void const* A, void const* B, + void const* input_sf, void const* weight_sf, + float const* global_sf, int m, int n, int k, + int batch_count, CutlassGemmConfig gemmConfig, + char* workspace, const size_t workspaceBytes, + cudaStream_t stream, int* occupancy = nullptr) { + // Several constraints: + // Cta N should be one of 128/192/256. + // M-mode size should be 128 or 256 for 2 CTA cluster MMA; + // M-mode size should be 128 for 1 CTA cluster OMMA. + // K256 looks to be better than K128 + switch (gemmConfig.tile_config_sm100) { + case CutlassTileConfigSM100::CtaShape128x64x128B: + return dispatchNVFP4xNVFP4GemmClusterShapeSm100, cute::Int<64>, + cute::Int<128>>( + D, A, B, input_sf, weight_sf, global_sf, m, n, k, batch_count, gemmConfig, workspace, + workspaceBytes, stream, occupancy); + break; + case CutlassTileConfigSM100::CtaShape128x256x128B: + return dispatchNVFP4xNVFP4GemmClusterShapeSm100, cute::Int<256>, + cute::Int<128>>( + D, A, B, input_sf, weight_sf, global_sf, m, n, k, batch_count, gemmConfig, workspace, + workspaceBytes, stream, occupancy); + break; + case CutlassTileConfigSM100::CtaShape128x128x256B: + return dispatchNVFP4xNVFP4GemmClusterShapeSm100, cute::Int<128>, + cute::Int<256>>( + D, A, B, input_sf, weight_sf, global_sf, m, n, k, batch_count, gemmConfig, workspace, + workspaceBytes, stream, occupancy); + break; + case CutlassTileConfigSM100::CtaShape128x256x256B: + return dispatchNVFP4xNVFP4GemmClusterShapeSm100, cute::Int<256>, + cute::Int<256>>( + D, A, B, input_sf, weight_sf, global_sf, m, n, k, batch_count, gemmConfig, workspace, + workspaceBytes, stream, occupancy); + break; + case CutlassTileConfigSM100::Undefined: + throw std::runtime_error("[Error][FP4][dispatch_gemm_cta_shape] Gemm config undefined."); + break; + case CutlassTileConfigSM100::ChooseWithHeuristic: + throw std::runtime_error( + "[Error][FP4][dispatch_gemm_cta_shape] Gemm config should have already been " + "set by " + "heuristic."); + break; + default: + throw std::runtime_error( + "[Error][FP4][dispatch_gemm_cta_shape] Config is invalid for FP4 GEMM."); + break; + } +} +template +CutlassFp4GemmRunner::CutlassFp4GemmRunner() {} + +template +CutlassFp4GemmRunner::~CutlassFp4GemmRunner() {} + +template +size_t CutlassFp4GemmRunner::dispatchToArch( + T* D, void const* A, void const* B, void const* input_sf, void const* weight_sf, + float const* global_sf, int m, int n, int k, int batch_count, CutlassGemmConfig gemmConfig, + char* workspace, const size_t workspaceBytes, cudaStream_t stream, int* occupancy) { + if constexpr (fp4GemmType == FP4GemmType::W4A4_NVFP4_NVFP4) { + return dispatchNVFP4xNVFP4GemmCTAShapeSm100(D, A, B, input_sf, weight_sf, global_sf, m, n, k, + batch_count, gemmConfig, workspace, + workspaceBytes, stream, occupancy); + } else { + throw std::runtime_error( + "[Error][CutlassFp4GemmRunner][GEMM Dispatch] FP4 Gemm type unsupported for " + "CUTLASS FP4 GEMM"); + } +} + +template +void CutlassFp4GemmRunner::gemm(void* D, void const* A, void const* B, + void const* input_sf, void const* weight_sf, + float const* global_sf, int m, int n, int k, + int batch_count, CutlassGemmConfig gemmConfig, + char* workspace, const size_t workspaceBytes, + cudaStream_t stream) { + CutlassFp4GemmRunner::dispatchToArch( + reinterpret_cast(D), A, B, input_sf, weight_sf, global_sf, m, n, k, batch_count, + gemmConfig, workspace, workspaceBytes, stream); +} + +template +std::vector CutlassFp4GemmRunner::getConfigs() const { + std::vector candidateConfigs; + + std::vector tilesSm100 = { + CutlassTileConfigSM100::CtaShape128x64x128B, + CutlassTileConfigSM100::CtaShape128x256x128B, + CutlassTileConfigSM100::CtaShape128x128x256B, + CutlassTileConfigSM100::CtaShape128x256x256B, + }; + std::vector clusterShapes = { + ClusterShape::ClusterShape_1x1x1, ClusterShape::ClusterShape_1x2x1, + ClusterShape::ClusterShape_2x1x1, ClusterShape::ClusterShape_2x2x1, + ClusterShape::ClusterShape_1x4x1, ClusterShape::ClusterShape_4x2x1, + ClusterShape::ClusterShape_2x4x1, ClusterShape::ClusterShape_4x4x1, + }; + for (auto const& tile_config : tilesSm100) { + for (auto const& cluster_config : clusterShapes) { + CutlassGemmConfig config(tile_config, MainloopScheduleType::AUTO, EpilogueScheduleType::AUTO, + cluster_config); + candidateConfigs.push_back(config); + } + } + + // There’s no heuristic yet, so for users without autotuning, we provide an ordering based on + // performance sweeps from common workloads. + std::vector best_tactics_index = {22, 20, 29, 4, 18}; + std::vector newCandidateConfigs; + for (auto const& tactic_index : best_tactics_index) { + newCandidateConfigs.push_back(candidateConfigs[tactic_index]); + } + for (int64_t i = 0; i < candidateConfigs.size(); i++) { + if (std::find(best_tactics_index.begin(), best_tactics_index.end(), i) == + best_tactics_index.end()) { + newCandidateConfigs.push_back(candidateConfigs[i]); + } + } + return newCandidateConfigs; +} + +template +size_t CutlassFp4GemmRunner::getWorkspaceSizeImpl(int const m, int const n, + int const k, + int const batch_count) { + size_t workspace_size = 0; + auto gemmConfigs = CutlassFp4GemmRunner{}.getConfigs(); + for (auto const& gemmConfig : gemmConfigs) { + try { + size_t curr_workspace_size = CutlassFp4GemmRunner::dispatchToArch( + nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, m, n, k, batch_count, gemmConfig, + nullptr, 0, 0); + workspace_size = std::max(workspace_size, curr_workspace_size); + } catch (std::runtime_error& e) { + // Swallow errors when SMEM exceeds maximum allowed + continue; + } + } + return workspace_size; +} + +template +size_t CutlassFp4GemmRunner::getWorkspaceSize(int const m, int const n, int const k, + int const batch_count) { + // Custom hash function for the MNKB type + using MNK = std::tuple; + + struct MNKHash { + size_t operator()(const MNK& mnk) const { + auto h1 = std::hash{}(std::get<0>(mnk)); + auto h2 = std::hash{}(std::get<1>(mnk)); + auto h3 = std::hash{}(std::get<2>(mnk)); + auto h4 = std::hash{}(std::get<3>(mnk)); + return h1 ^ h2 ^ h3 ^ h4; + } + }; + + static std::unordered_map workspace_hashmap; + + size_t workspace_size = 0; + if (workspace_hashmap.find(std::make_tuple(m, n, k, batch_count)) == workspace_hashmap.end()) { + workspace_size = + CutlassFp4GemmRunner::getWorkspaceSizeImpl(m, n, k, batch_count); + workspace_hashmap[std::make_tuple(m, n, k, batch_count)] = workspace_size; + } else { + workspace_size = workspace_hashmap[std::make_tuple(m, n, k, batch_count)]; + } + return workspace_size; +} + +} // namespace gemm +} // namespace flashinfer +#endif // FLASHINFER_FP4_GEMM_CUTLASS_TEMPLATE_H_ diff --git a/include/flashinfer/gemm/fp4_gemm_template_sm100.h b/include/flashinfer/gemm/fp4_gemm_template_sm100.h new file mode 100644 index 0000000000000000000000000000000000000000..62548686c793a397364796b5fb1159db268c2140 --- /dev/null +++ b/include/flashinfer/gemm/fp4_gemm_template_sm100.h @@ -0,0 +1,288 @@ +/* + * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef FLASHINFER_FP4_GEMM_TEMPLATE_SM100_H_ +#define FLASHINFER_FP4_GEMM_TEMPLATE_SM100_H_ + +#ifndef _WIN32 +#pragma GCC diagnostic push +#pragma GCC diagnostic ignored "-Wstrict-aliasing" +#endif // #ifndef _WIN32 + +#include "cutlass/arch/arch.h" +#include "cutlass/cutlass.h" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/gemm.h" +#include "flashinfer/arch_condition.h" +#include "flashinfer/cutlass_utils.cuh" + +#ifndef _WIN32 +#pragma GCC diagnostic pop +#endif // #ifndef _WIN32 + +namespace flashinfer { +namespace gemm { +using namespace cute; + +#ifdef ENABLE_BF16 +using SafeBF16 = __nv_bfloat16; +#else +using SafeBF16 = void; +#endif + +struct _1SM {}; + +struct _2SM {}; + +template +struct SMTypeAdapter {}; + +template <> +struct SMTypeAdapter<_1SM> { + static int const Scale = 1; + using AtomThrShape = cute::Shape<_1, _1, _1>; + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized1Sm; + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized1SmNvf4Sm100; +}; + +template <> +struct SMTypeAdapter<_2SM> { + static int const Scale = 2; + using AtomThrShape = cute::Shape<_2, _1, _1>; + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized2Sm; + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized2SmNvf4Sm100; +}; + +template +constexpr auto always_false = false; + +template +size_t genericFp4GemmKernelLauncher(void* D, void const* A, void const* B, void const* input_sf, + void const* weight_sf, float const* global_sf, int m, int n, + int k, int batch_count, CutlassGemmConfig gemmConfig, + char* workspace, size_t const workspaceBytes, + cudaStream_t stream, int* occupancy); + +#ifdef PLACEHOLDER_KERNELS + +#define INSTANTIATE_FP4_GEMM_KERNEL_LAUNCHER(T, CTA_M_, CTA_N_, CTA_K_, CGA_M_, CGA_N_, CGA_K_, \ + XSM_) \ + template <> \ + size_t \ + genericFp4GemmKernelLauncher, cute::Int, cute::Int, \ + cute::Int, cute::Int, cute::Int, XSM_>( \ + void* D, void const* A, void const* B, void const* input_sf, void const* weight_sf, \ + float const* global_sf, int m, int n, int k, int batch_count, CutlassGemmConfig gemmConfig, \ + char* workspace, const size_t workspaceBytes, cudaStream_t stream, int* occupancy) { \ + throw std::runtime_error( \ + "FP4 gemm kernel is not compiled with support for " \ + "this Architecture."); \ + } + +#else + +#define INSTANTIATE_FP4_GEMM_KERNEL_LAUNCHER(T, CTA_M_, CTA_N_, CTA_K_, CGA_M_, CGA_N_, CGA_K_, \ + XSM_) \ + struct \ + DeviceGemmFp4GemmSm100_##T##_##CTA_M_##_##CTA_N_##_##CTA_K_##_##CGA_M_##_##CGA_N_##_##CGA_K_##XSM_ { \ + using OutElementType = flashinfer::cutlass_dtype::type; \ + using CTAShape = cute::Shape, cute::Int, cute::Int>; \ + /*using ClusterShape = cute::Shape, cute::Int, cute::Int>;*/ \ + using ClusterShape = cute::Shape; \ + using ElementType = cutlass::float_e2m1_t; \ + using Arch = cutlass::arch::Sm100; \ + /* // Input A */ \ + using ElementA = ElementType; \ + using LayoutA = cutlass::layout::RowMajor; \ + static constexpr int AlignmentA = 128 / cutlass::sizeof_bits::value; \ + /* // Input B */ \ + using ElementB = ElementType; \ + using LayoutB = cutlass::layout::ColumnMajor; \ + static constexpr int AlignmentB = 128 / cutlass::sizeof_bits::value; \ + /* // Input C */ \ + using ElementC = void; \ + using LayoutC = cutlass::layout::RowMajor; \ + static constexpr int AlignmentC = 128 / cutlass::sizeof_bits::value; \ + \ + using SFType = cutlass::float_ue4m3_t; \ + using ElementCompute = float; \ + using ElementAccumulator = float; \ + using OperatorClass = cutlass::arch::OpClassTensorOp; \ + using EpilogueTileType = std::conditional_t, \ + cutlass::epilogue::collective::EpilogueTileAuto>; \ + using EpilogueSchedule = SMTypeAdapter::EpilogueSchedule; \ + using MainloopSchedule = SMTypeAdapter::MainloopSchedule; \ + using MmaTileShape = cute::Shape::Scale>, \ + cute::Int, cute::Int>; \ + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< \ + Arch, OperatorClass, MmaTileShape, ClusterShape, EpilogueTileType, ElementAccumulator, \ + ElementCompute, ElementC, LayoutC, AlignmentC, OutElementType, LayoutC, AlignmentC, \ + EpilogueSchedule, \ + cutlass::epilogue::fusion::LinearCombination>::CollectiveOp; \ + \ + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< \ + Arch, cutlass::arch::OpClassBlockScaledTensorOp, cute::tuple, LayoutA, \ + AlignmentA, cute::tuple, LayoutB, AlignmentB, ElementAccumulator, \ + MmaTileShape, ClusterShape, \ + cutlass::gemm::collective::StageCountAutoCarveout( \ + sizeof(typename CollectiveEpilogue::SharedStorage))>, \ + MainloopSchedule>::CollectiveOp; \ + \ + template \ + struct Sm10xOnly : Base { \ + using typename Base::Params; \ + CUTLASS_DEVICE \ + void operator()(Params const& params, char* smem_buf) { \ + if constexpr (flashinfer::arch::is_major_v<10>) { \ + this->Base::operator()(params, smem_buf); \ + } else { \ + if (cute::thread0()) { \ + printf("%s : This kernel shall only run on SM10x devices.\n", __PRETTY_FUNCTION__); \ + __trap(); \ + } \ + } \ + } \ + }; \ + using GemmKernel = \ + Sm10xOnly, \ + CollectiveMainloop, CollectiveEpilogue, \ + cutlass::gemm::PersistentScheduler>>; \ + \ + using Gemm = typename cutlass::gemm::device::GemmUniversalAdapter; \ + }; \ + \ + template \ + typename Gemm::Arguments \ + prepareGemmArgs_##T##_##CTA_M_##_##CTA_N_##_##CTA_K_##_##CGA_M_##_##CGA_N_##_##CGA_K_##XSM_( \ + void* D, void const* A, void const* B, void const* input_sf, void const* weight_sf, \ + float const* global_sf, int m, int n, int k, int batch_count) { \ + using Sm1xxBlkScaledConfig = \ + typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig; \ + using ElementA = typename Gemm::ElementA; \ + using ElementB = typename Gemm::ElementB; \ + using ElementSFA = cutlass::float_ue4m3_t; \ + using ElementSFB = cutlass::float_ue4m3_t; \ + using ElementC = void; \ + using ElementD = typename Gemm::ElementD; \ + using ElementCompute = float; \ + \ + typename Gemm::Arguments operator_args; \ + operator_args.mode = cutlass::gemm::GemmUniversalMode::kGemm; \ + auto& fusion_args = operator_args.epilogue.thread; \ + fusion_args.alpha_ptr = static_cast(global_sf); \ + \ + operator_args.problem_shape = cute::make_shape(m, n, k, batch_count); \ + \ + operator_args.mainloop.ptr_A = static_cast(A); \ + operator_args.mainloop.ptr_B = static_cast(B); \ + operator_args.mainloop.ptr_SFA = static_cast(input_sf); \ + operator_args.mainloop.ptr_SFB = static_cast(weight_sf); \ + operator_args.epilogue.ptr_C = static_cast(D); \ + operator_args.epilogue.ptr_D = static_cast(D); \ + \ + int const stride_A = batch_count == 1 ? 0 : m * k; \ + int const stride_B = batch_count == 1 ? 0 : n * k; \ + int const stride_C = batch_count == 1 ? 0 : m * n; \ + \ + operator_args.mainloop.dA = \ + cute::make_int_tuple_from(k, stride_A); \ + operator_args.mainloop.dB = \ + cute::make_int_tuple_from(k, stride_B); \ + operator_args.epilogue.dC = \ + cute::make_int_tuple_from(n, stride_C); \ + operator_args.epilogue.dD = operator_args.epilogue.dC; \ + \ + operator_args.mainloop.layout_SFA = \ + Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(operator_args.problem_shape); \ + operator_args.mainloop.layout_SFB = \ + Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(operator_args.problem_shape); \ + \ + if constexpr (!std::is_const_v) { \ + operator_args.scheduler.max_swizzle_size = 1; \ + } \ + if constexpr (!std::is_const_v) { \ + using Enum_t = decltype(operator_args.scheduler.raster_order); \ + operator_args.scheduler.raster_order = Enum_t::Heuristic; \ + } \ + operator_args.hw_info.cluster_shape = dim3(CGA_M_, CGA_N_, CGA_K_); \ + operator_args.hw_info.cluster_shape_fallback = dim3(SMTypeAdapter::Scale, 1, 1); \ + \ + return operator_args; \ + } \ + \ + template <> \ + size_t \ + genericFp4GemmKernelLauncher, cute::Int, cute::Int, \ + cute::Int, cute::Int, cute::Int, XSM_>( \ + void* D, void const* A, void const* B, void const* input_sf, void const* weight_sf, \ + float const* global_sf, int m, int n, int k, int batch_count, CutlassGemmConfig gemmConfig, \ + char* workspace, const size_t workspaceBytes, cudaStream_t stream, int* occupancy) { \ + using ElementOutput__ = \ + typename cutlass::platform::conditional::value, \ + cutlass::half_t, T>::type; \ + using ElementOutput_ = typename cutlass::platform::conditional< \ + cutlass::platform::is_same::value, float, ElementOutput__>::type; \ + using ElementOutput = typename cutlass::platform::conditional< \ + cutlass::platform::is_same::value, cutlass::bfloat16_t, \ + ElementOutput_>::type; \ + \ + using Fp4GemmOperator = \ + DeviceGemmFp4GemmSm100_##T##_##CTA_M_##_##CTA_N_##_##CTA_K_##_##CGA_M_##_##CGA_N_##_##CGA_K_##XSM_:: \ + Gemm; \ + Fp4GemmOperator gemm; \ + auto args = \ + prepareGemmArgs_##T##_##CTA_M_##_##CTA_N_##_##CTA_K_##_##CGA_M_##_##CGA_N_##_##CGA_K_##XSM_< \ + Fp4GemmOperator>(D, A, B, input_sf, weight_sf, global_sf, m, n, k, batch_count); \ + /* // Return workspace size */ \ + if (!A && !B && !D) { \ + return gemm.get_workspace_size(args); \ + } \ + if (gemm.get_workspace_size(args) > workspaceBytes) { \ + std::string errMsg("Requested workspace size insufficient. Required " + \ + std::to_string(gemm.get_workspace_size(args)) + ", got " + \ + std::to_string(workspaceBytes)); \ + throw std::runtime_error("[FP4 gemm Runner] " + errMsg); \ + } \ + auto can_implement = gemm.can_implement(args); \ + if (can_implement != cutlass::Status::kSuccess) { \ + std::string errMsg = "FP4 Gemm cutlass kernel will fail for params. Error: " + \ + std::string(cutlassGetStatusString(can_implement)); \ + throw std::runtime_error("[FP4 gemm Runner] " + errMsg); \ + } \ + auto initStatus = gemm.initialize(args, workspace, stream); \ + if (initStatus != cutlass::Status::kSuccess) { \ + std::string errMsg = "Failed to initialize cutlass FP4 gemm. Error: " + \ + std::string(cutlassGetStatusString(initStatus)); \ + throw std::runtime_error("[FP4 gemm Runner] " + errMsg); \ + } \ + auto runStatus = gemm.run(args, workspace, stream, nullptr, /* enablePDL */ false); \ + if (runStatus != cutlass::Status::kSuccess) { \ + std::string errMsg = "Failed to run cutlass FP4 gemm. Error: " + \ + std::string(cutlassGetStatusString(runStatus)); \ + throw std::runtime_error("[FP4 gemm Runner] " + errMsg); \ + } \ + return gemm.get_workspace_size(args); \ + } + +#endif + +} // namespace gemm +} // namespace flashinfer +#endif // FLASHINFER_FP4_GEMM_TEMPLATE_SM100_H_ diff --git a/include/flashinfer/gemm/fp8_gemm_cutlass.h b/include/flashinfer/gemm/fp8_gemm_cutlass.h new file mode 100644 index 0000000000000000000000000000000000000000..3168b66b2a3f202e9d437e28b498ab6de81e146a --- /dev/null +++ b/include/flashinfer/gemm/fp8_gemm_cutlass.h @@ -0,0 +1,62 @@ +/* + * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#ifndef FLASHINFER_FP8_GEMM_CUTLASS_H_ +#define FLASHINFER_FP8_GEMM_CUTLASS_H_ + +#include + +#include + +#include "flashinfer/gemm/cutlass_gemm_configs.h" + +namespace flashinfer { +namespace gemm { + +class CutlassFp8GemmRunnerInterface { + public: + CutlassFp8GemmRunnerInterface() = default; + virtual ~CutlassFp8GemmRunnerInterface() = default; + + virtual void gemm(__nv_fp8_e4m3 const* A, __nv_fp8_e4m3 const* B, float const* alpha, void* D, + int m, int n, int k, int b, CutlassGemmConfig gemmConfig, char* workspacePtr, + size_t const workspaceBytes, cudaStream_t stream) = 0; + + virtual size_t getWorkspaceSize(int m, int n, int k) = 0; + + virtual std::vector getConfigs() const = 0; +}; + +template +class CutlassFp8GemmRunner : public virtual CutlassFp8GemmRunnerInterface { + public: + CutlassFp8GemmRunner() = default; + ~CutlassFp8GemmRunner() = default; + + void gemm(__nv_fp8_e4m3 const* A, __nv_fp8_e4m3 const* B, float const* alpha, void* D, int m, + int n, int k, int b, CutlassGemmConfig gemmConfig, char* workspacePtr, + size_t const workspaceBytes, cudaStream_t stream) override; + size_t getWorkspaceSize(int m, int n, int k) override; + std::vector getConfigs() const override; + + private: + size_t getWorkspaceSizeImpl(int m, int n, int k); +}; + +} // namespace gemm +} // namespace flashinfer + +#endif // FLASHINFER_FP8_GEMM_CUTLASS_H_ diff --git a/include/flashinfer/gemm/fp8_gemm_cutlass_template.h b/include/flashinfer/gemm/fp8_gemm_cutlass_template.h new file mode 100644 index 0000000000000000000000000000000000000000..b9dc2f738b7dce27b7afbf4637ddb15136be52a8 --- /dev/null +++ b/include/flashinfer/gemm/fp8_gemm_cutlass_template.h @@ -0,0 +1,227 @@ +/* + * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef FLASHINFER_FP8_GEMM_CUTLASS_TEMPLATE_H_ +#define FLASHINFER_FP8_GEMM_CUTLASS_TEMPLATE_H_ + +#ifdef __GNUC__ // Check if the compiler is GCC or Clang +#pragma GCC diagnostic push +#pragma GCC diagnostic ignored "-Wstrict-aliasing" +#endif // __GNUC__ +#include "cutlass/arch/arch.h" +#include "cutlass/cutlass.h" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/gemm.h" +#include "flashinfer/arch_condition.h" +#include "flashinfer/cutlass_utils.cuh" + +#ifdef __GNUC__ // Check if the compiler is GCC or Clang +#pragma GCC diagnostic pop +#endif // __GNUC__ + +#include +#include + +#include "cutlass/float8.h" + +namespace flashinfer { +namespace gemm { + +struct _1SM {}; +struct _2SM {}; + +// only support fp8 TN gemm +// batch GEMM +// return require_workspace_size +template +size_t genericFp8GemmKernelLauncherSm100(__nv_fp8_e4m3 const* A, __nv_fp8_e4m3 const* B, + float const* alpha, T* D, int m, int n, int k, int b, + CutlassGemmConfig config, char* workspacePtr, + size_t const workspaceBytes, cudaStream_t stream); + +template +size_t dispatchGemmClusterShapeSm100(__nv_fp8_e4m3 const* A, __nv_fp8_e4m3 const* B, + float const* alpha, T* D, int m, int n, int k, int b, + CutlassGemmConfig gemmConfig, char* workspacePtr, + size_t const workspaceBytes, cudaStream_t stream) { + using namespace cute; + + switch (gemmConfig.cluster_shape) { + case ClusterShape::ClusterShape_1x1x1: + return genericFp8GemmKernelLauncherSm100, + _1SM>(A, B, alpha, D, m, n, k, b, gemmConfig, + workspacePtr, workspaceBytes, stream); + break; + + case ClusterShape::ClusterShape_2x1x1: + return genericFp8GemmKernelLauncherSm100, + _2SM>(A, B, alpha, D, m, n, k, b, gemmConfig, + workspacePtr, workspaceBytes, stream); + break; + case ClusterShape::ClusterShape_1x2x1: + return genericFp8GemmKernelLauncherSm100, + _1SM>(A, B, alpha, D, m, n, k, b, gemmConfig, + workspacePtr, workspaceBytes, stream); + break; + case ClusterShape::ClusterShape_2x2x1: + return genericFp8GemmKernelLauncherSm100, + _2SM>(A, B, alpha, D, m, n, k, b, gemmConfig, + workspacePtr, workspaceBytes, stream); + break; + default: + throw std::runtime_error("invalid config for fp8 gemm"); + break; + } +} + +template +size_t dispatchToArch(__nv_fp8_e4m3 const* A, __nv_fp8_e4m3 const* B, float const* alpha, void* D, + int m, int n, int k, int b, CutlassGemmConfig gemmConfig, char* workspacePtr, + size_t const workspaceBytes, cudaStream_t stream) { + using namespace cute; + + using arch = cutlass::arch::Sm100; + + // swap M, N and A,B + // change m-> n , n-> m + // A rowmajor, B colmajor , C, D rowmajor + switch (gemmConfig.tile_config_sm100) { + case CutlassTileConfigSM100::CtaShape64x64x128B: + return dispatchGemmClusterShapeSm100(B, A, alpha, static_cast(D), n, + m, k, b, gemmConfig, workspacePtr, + workspaceBytes, stream); + break; + case CutlassTileConfigSM100::CtaShape64x128x128B: + return dispatchGemmClusterShapeSm100( + B, A, alpha, static_cast(D), n, m, k, b, gemmConfig, workspacePtr, workspaceBytes, + stream); + break; + case CutlassTileConfigSM100::CtaShape64x256x128B: + return dispatchGemmClusterShapeSm100( + B, A, alpha, static_cast(D), n, m, k, b, gemmConfig, workspacePtr, workspaceBytes, + stream); + break; + case CutlassTileConfigSM100::CtaShape128x64x128B: + return dispatchGemmClusterShapeSm100( + B, A, alpha, static_cast(D), n, m, k, b, gemmConfig, workspacePtr, workspaceBytes, + stream); + break; + case CutlassTileConfigSM100::CtaShape128x128x128B: + return dispatchGemmClusterShapeSm100( + B, A, alpha, static_cast(D), n, m, k, b, gemmConfig, workspacePtr, workspaceBytes, + stream); + break; + case CutlassTileConfigSM100::CtaShape128x256x128B: + return dispatchGemmClusterShapeSm100( + B, A, alpha, static_cast(D), n, m, k, b, gemmConfig, workspacePtr, workspaceBytes, + stream); + break; + + default: + throw std::runtime_error("unsupported tile config for fp8 gemm"); + break; + } +} + +template +void CutlassFp8GemmRunner::gemm(__nv_fp8_e4m3 const* A, __nv_fp8_e4m3 const* B, + float const* alpha, void* D, int m, int n, int k, int b, + CutlassGemmConfig gemmConfig, char* workspacePtr, + size_t const workspaceBytes, cudaStream_t stream) { + dispatchToArch(A, B, alpha, reinterpret_cast(D), m, n, k, b, gemmConfig, workspacePtr, + workspaceBytes, stream); +} + +template +size_t CutlassFp8GemmRunner::getWorkspaceSizeImpl(int m, int n, int k) { + size_t workspace_size = 0; + auto gemmConfigs = CutlassFp8GemmRunner{}.getConfigs(); + for (auto const& gemmConfig : gemmConfigs) { + try { + size_t curr_workspace_size = dispatchToArch(nullptr, nullptr, nullptr, nullptr, m, n, k, 1, + gemmConfig, nullptr, 0, nullptr); + + workspace_size = std::max(workspace_size, curr_workspace_size); + } catch (std::runtime_error& e) { + // Swallow errors when SMEM exceeds maximum allowed + continue; + } + } + + return workspace_size; +} + +template +size_t CutlassFp8GemmRunner::getWorkspaceSize(int m, int n, int k) { + // Custom hash function for the MNK type + using MNK = std::tuple; + + struct MNKHash { + size_t operator()(const MNK& mnk) const { + auto h1 = std::hash{}(std::get<0>(mnk)); + auto h2 = std::hash{}(std::get<1>(mnk)); + auto h3 = std::hash{}(std::get<2>(mnk)); + return h1 ^ h2 ^ h3; + } + }; + + static std::unordered_map workspace_hashmap; + + size_t workspace_size = 0; + if (workspace_hashmap.find(std::make_tuple(m, n, k)) == workspace_hashmap.end()) { + workspace_size = CutlassFp8GemmRunner::getWorkspaceSizeImpl(m, n, k); + workspace_hashmap[std::make_tuple(m, n, k)] = workspace_size; + } else { + workspace_size = workspace_hashmap[std::make_tuple(m, n, k)]; + } + return workspace_size; +} + +template +std::vector CutlassFp8GemmRunner::getConfigs() const { + std::vector candidate_configs; + + std::vector candidateConfigs; + + std::vector tilesSm100 = { + CutlassTileConfigSM100::CtaShape64x64x128B, CutlassTileConfigSM100::CtaShape64x128x128B, + CutlassTileConfigSM100::CtaShape64x256x128B, CutlassTileConfigSM100::CtaShape128x64x128B, + CutlassTileConfigSM100::CtaShape128x128x128B, CutlassTileConfigSM100::CtaShape128x256x128B, + }; + + std::vector clusterShapes = { + ClusterShape::ClusterShape_1x1x1, + ClusterShape::ClusterShape_1x2x1, + ClusterShape::ClusterShape_2x1x1, + ClusterShape::ClusterShape_2x2x1, + }; + for (auto const& tile_config : tilesSm100) { + for (auto const& cluster_config : clusterShapes) { + CutlassGemmConfig config(tile_config, MainloopScheduleType::AUTO, EpilogueScheduleType::AUTO, + cluster_config); + candidateConfigs.push_back(config); + } + } + + return candidateConfigs; +} + +} // namespace gemm +} // namespace flashinfer + +#endif // FLASHINFER_FP8_GEMM_CUTLASS_TEMPLATE_H_ diff --git a/include/flashinfer/gemm/fp8_gemm_template_sm100.h b/include/flashinfer/gemm/fp8_gemm_template_sm100.h new file mode 100644 index 0000000000000000000000000000000000000000..fd35975cddb9c3bfa6466d4db9d3f377d5186487 --- /dev/null +++ b/include/flashinfer/gemm/fp8_gemm_template_sm100.h @@ -0,0 +1,217 @@ +/* + * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef FLASHINFER_FP8_GEMM_TEMPLATE_SM100_H_ +#define FLASHINFER_FP8_GEMM_TEMPLATE_SM100_H_ + +#ifdef __GNUC__ // Check if the compiler is GCC or Clang +#pragma GCC diagnostic push +#pragma GCC diagnostic ignored "-Wstrict-aliasing" +#endif // __GNUC__ +#include "cutlass/arch/arch.h" +#include "cutlass/cutlass.h" +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/gemm/collective/collective_builder.hpp" +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/gemm.h" +#include "flashinfer/arch_condition.h" +#include "flashinfer/cutlass_utils.cuh" + +#ifdef __GNUC__ // Check if the compiler is GCC or Clang +#pragma GCC diagnostic pop +#endif // __GNUC__ + +#include +#include + +#include "cutlass/float8.h" +#include "flashinfer/gemm/cutlass_gemm_configs.h" + +namespace flashinfer { +namespace gemm { + +template +struct SMTypeAdapter {}; + +struct _1SM; +struct _2SM; + +template <> +struct SMTypeAdapter<_1SM> { + static int const Scale = 1; + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized1Sm; + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized1SmSm100; +}; + +template <> +struct SMTypeAdapter<_2SM> { + static int const Scale = 2; + using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized2Sm; + using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized2SmSm100; +}; + +// only support fp8 TN gemm +// batch GEMM +// return require_workspace_size +template +size_t genericFp8GemmKernelLauncherSm100(__nv_fp8_e4m3 const* A, __nv_fp8_e4m3 const* B, + float const* alpha, T* D, int m, int n, int k, int b, + CutlassGemmConfig config, char* workspacePtr, + size_t const workspaceBytes, cudaStream_t stream) { + using namespace cute; + + // A matrix configuration + using ElementA = cutlass::float_e4m3_t; // Element type for A matrix operand + using LayoutA = cutlass::layout::RowMajor; // Layout type for A matrix operand + constexpr int AlignmentA = + 128 / cutlass::sizeof_bits::value; // Memory access granularity/alignment of A + // matrix in units of elements (up to 16 bytes) + + // B matrix configuration + using ElementB = cutlass::float_e4m3_t; // Element type for B matrix operand + using LayoutB = cutlass::layout::ColumnMajor; // Layout type for B matrix operand + constexpr int AlignmentB = + 128 / cutlass::sizeof_bits::value; // Memory access granularity/alignment of B + // matrix in units of elements (up to 16 bytes) + + using ElementOutput_ = + typename cutlass::platform::conditional::value, + cutlass::half_t, T>::type; +#ifdef ENABLE_BF16 + using ElementOutput = typename cutlass::platform::conditional< + cutlass::platform::is_same::value, cutlass::bfloat16_t, + ElementOutput_>::type; +#else + using ElementOutput = ElementOutput_; +#endif + + // C matrix configuration + using ElementC = ElementOutput; // Element type for C and D matrix operands + using LayoutC = cutlass::layout::ColumnMajor; // Layout type for C and D matrix operands + constexpr int AlignmentC = + 128 / cutlass::sizeof_bits::value; // Memory access granularity/alignment of C + // matrix in units of elements (up to 16 bytes) + + // D matrix configuration + using ElementD = ElementC; + using LayoutD = LayoutC; + constexpr int AlignmentD = AlignmentC; + + // / Core kernel configurations + using ElementAccumulator = float; // Element type for internal accumulation + using ElementCompute = float; // Element type for epilogue computation + using ArchTag = + cutlass::arch::Sm100; // Tag indicating the minimum SM that supports the intended feature + using OperatorClass = cutlass::arch::OpClassTensorOp; // Operator class tag + using TileShape = cute::Shape::Scale>, cute::Int, + cute::Int>; // Threadblock-level tile size + + using ClusterShape = ClusterShape_; // Shape of the threadblocks in a cluster + using EpilogueSchedule = typename SMTypeAdapter::EpilogueSchedule; + using MainloopSchedule = typename SMTypeAdapter::MainloopSchedule; + using EpilogueTileType = cutlass::epilogue::collective::EpilogueTileAuto; + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + ArchTag, OperatorClass, TileShape, ClusterShape, EpilogueTileType, ElementAccumulator, + ElementCompute, ElementC, LayoutC, AlignmentC, ElementD, LayoutD, AlignmentD, + EpilogueSchedule>::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + ArchTag, OperatorClass, ElementA, LayoutA, AlignmentA, ElementB, LayoutB, AlignmentB, + ElementAccumulator, TileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout( + sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule>::CollectiveOp; + + using GemmKernel = + cutlass::gemm::kernel::GemmUniversal, // Indicates ProblemShape + CollectiveMainloop, CollectiveEpilogue>; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + + // Extract information from Gemm kernel. + using EpilogueOutputOp = typename Gemm::EpilogueOutputOp; + using ElementScalar = typename EpilogueOutputOp::ElementScalar; + + using StrideA = typename Gemm::GemmKernel::StrideA; + using StrideB = typename Gemm::GemmKernel::StrideB; + using StrideC = typename Gemm::GemmKernel::StrideC; + using StrideD = typename Gemm::GemmKernel::StrideD; + + using LayoutScalar = cutlass::layout::PackedVectorLayout; + + StrideA stride_A = cutlass::make_cute_packed_stride(StrideA{}, cute::make_shape(m, k, b)); + StrideB stride_B = cutlass::make_cute_packed_stride(StrideB{}, cute::make_shape(n, k, b)); + StrideC stride_C = cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(m, n, b)); + StrideD stride_D = cutlass::make_cute_packed_stride(StrideD{}, cute::make_shape(m, n, b)); + + typename Gemm::Arguments arguments{cutlass::gemm::GemmUniversalMode::kGemm, + {m, n, k, b}, + {reinterpret_cast(A), stride_A, + reinterpret_cast(B), stride_B}, + {{}, // epilogue.thread + nullptr, + stride_C, + reinterpret_cast(D), + stride_D}}; + + auto& fusion_args = arguments.epilogue.thread; + fusion_args.alpha = 0.F; + fusion_args.beta = 0.F; + fusion_args.alpha_ptr = alpha; + fusion_args.beta_ptr = nullptr; + + Gemm gemm; + + // Return workspace size + if (!A && !B && !D) { + return Gemm::get_workspace_size(arguments); + } + + if (Gemm::get_workspace_size(arguments) > workspaceBytes) { + throw std::runtime_error("[Fp8 Gemm Runner] failed to get workspace size"); + } + + auto can_implement = gemm.can_implement(arguments); + if (can_implement != cutlass::Status::kSuccess) { + throw std::runtime_error("[Fp8 Gemm Runner] cutlass kernel not implemented given the params"); + } + + auto initStatus = gemm.initialize(arguments, workspacePtr); + if (initStatus != cutlass::Status::kSuccess) { + throw std::runtime_error("[Fp8 Gemm Runner] failed to initialize"); + } + + auto runStatus = gemm.run(stream, nullptr, true); + if (runStatus != cutlass::Status::kSuccess) { + throw std::runtime_error("[Fp8 Gemm Runner] failed to run"); + } + return gemm.get_workspace_size(arguments); +} + +} // namespace gemm +} // namespace flashinfer + +#define INSTANCE_FP8_GEMM_TEMPLATE_SM100(RET_TYPE, TILE_M, TILE_N, TILE_K, CGA_M_, CGA_N_, CGA_K_, \ + SM_TYPE) \ + template size_t genericFp8GemmKernelLauncherSm100< \ + RET_TYPE, cutlass::arch::Sm100, TILE_M, TILE_N, TILE_K, \ + cute::Shape, cute::Int, cute::Int>, SM_TYPE>( \ + __nv_fp8_e4m3 const* A, __nv_fp8_e4m3 const* B, float const* alpha, RET_TYPE* D, int m, \ + int n, int k, int b, CutlassGemmConfig config, char* workspacePtr, \ + size_t const workspaceBytes, cudaStream_t stream); + +#endif // FLASHINFER_FP8_GEMM_TEMPLATE_SM100_H_ diff --git a/include/flashinfer/gemm/gemm_groupwise_sm100.cuh b/include/flashinfer/gemm/gemm_groupwise_sm100.cuh new file mode 100644 index 0000000000000000000000000000000000000000..a5ae6a6f73940243d844436195ff90d051d454ea --- /dev/null +++ b/include/flashinfer/gemm/gemm_groupwise_sm100.cuh @@ -0,0 +1,155 @@ +/* + * Copyright (c) 2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef FLASHINFER_GEMM_GROUPWISE_SM100_CUH_ +#define FLASHINFER_GEMM_GROUPWISE_SM100_CUH_ + +#include + +#include "../allocator.h" +#include "../cutlass_utils.cuh" +#include "../utils.cuh" + +namespace flashinfer { + +namespace gemm { + +using namespace cute; + +template +cudaError_t CutlassGroupwiseScaledGEMMSM100(void* float_buffer, size_t float_buffer_size_in_bytes, + DTypeIn* A_ptr, DTypeIn* B_ptr, float* SFA_ptr, + float* SFB_ptr, DTypeOut* C_ptr, int m, int n, int k, + int l, cudaStream_t stream) { + using ElementA = DTypeIn; // Element type for A matrix operand + using LayoutA = cutlass::layout::RowMajor; // Layout type for A matrix operand + constexpr int AlignmentA = + 128 / cutlass::sizeof_bits::value; // Memory access granularity/alignment of A + // matrix in units of elements (up to 16 bytes) + + // B matrix configuration + using ElementB = DTypeIn; // Element type for B matrix operand + using LayoutB = cutlass::layout::ColumnMajor; // Layout type for B matrix operand + constexpr int AlignmentB = + 128 / cutlass::sizeof_bits::value; // Memory access granularity/alignment of A + // matrix in units of elements (up to 16 bytes) + + // C/D matrix configuration + using ElementC = DTypeOut; // Element type for C and D matrix operands + using LayoutC = cutlass::layout::RowMajor; // Layout type for C and D matrix operands + constexpr int AlignmentC = + 128 / cutlass::sizeof_bits::value; // Memory access granularity/alignment of A + // matrix in units of elements (up to 16 bytes) + + using ElementD = ElementC; + using LayoutD = LayoutC; + constexpr int AlignmentD = AlignmentC; + + // MMA type + using ElementAccumulator = float; // Element Accumulator will also be our scale factor type + using ElementCompute = float; + + using MmaTileShape_MNK = Shape, _128, _128>; + using ClusterShape_MNK = Shape, _1, _1>; + + // NOTE(Zihao):: UMMA::Major::MN, UMMA::Major::MN is the fastest configuration. + + using ScaleConfig = std::conditional_t< + ScaleMajorK, + cutlass::detail::Sm100BlockwiseScaleConfig, + cutlass::detail::Sm100BlockwiseScaleConfig>; + + using LayoutSFA = + decltype(ScaleConfig::deduce_layoutSFA()); // Layout type for SFA matrix operand + using LayoutSFB = + decltype(ScaleConfig::deduce_layoutSFB()); // Layout type for SFB matrix operand + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, MmaTileShape_MNK, ClusterShape_MNK, + cutlass::epilogue::collective::EpilogueTileAuto, ElementAccumulator, ElementCompute, ElementC, + LayoutC, AlignmentC, ElementD, LayoutC, AlignmentD, + cutlass::epilogue::collective::EpilogueScheduleAuto>::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, ElementA, + cute::tuple, AlignmentA, ElementB, cute::tuple, + AlignmentB, ElementAccumulator, MmaTileShape_MNK, ClusterShape_MNK, + cutlass::gemm::collective::StageCountAutoCarveout( + sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::KernelScheduleSm100Blockwise>::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, CollectiveMainloop, CollectiveEpilogue, + void>; // Default to ClusterLaunchControl (CLC) based tile scheduler + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + + using StrideA = typename Gemm::GemmKernel::StrideA; + using StrideB = typename Gemm::GemmKernel::StrideB; + using StrideC = typename Gemm::GemmKernel::StrideC; + using StrideD = typename Gemm::GemmKernel::StrideD; + + auto stride_A = cutlass::make_cute_packed_stride(StrideA{}, cute::make_shape(m, k, l)); + auto stride_B = cutlass::make_cute_packed_stride(StrideB{}, cute::make_shape(n, k, l)); + auto stride_C = cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(m, n, l)); + auto stride_D = cutlass::make_cute_packed_stride(StrideD{}, cute::make_shape(m, n, l)); + + auto layout_SFA = ScaleConfig::tile_atom_to_shape_SFA(make_shape(m, n, k, l)); + auto layout_SFB = ScaleConfig::tile_atom_to_shape_SFB(make_shape(m, n, k, l)); + + typename Gemm::Arguments arguments{cutlass::gemm::GemmUniversalMode::kGemm, + {m, n, k, l}, + { + A_ptr, + stride_A, + B_ptr, + stride_B, + SFA_ptr, + layout_SFA, + SFB_ptr, + layout_SFB, + }, + { + {}, // epilogue.thread + C_ptr, + stride_C, + C_ptr, + stride_C, + }}; + auto& fusion_args = arguments.epilogue.thread; + fusion_args.alpha = 1.0f; + fusion_args.beta = 0.0f; + + Gemm gemm; + + size_t workspace_size = Gemm::get_workspace_size(arguments); + AlignedAllocator float_allocator(float_buffer, float_buffer_size_in_bytes); + auto workspace_ptr = float_allocator.aligned_alloc(workspace_size, 16, + "sm100_groupwise_gemm_float_workspace"); + + CUTLASS_CHECK(gemm.can_implement(arguments)); + CUTLASS_CHECK(gemm.initialize(arguments, workspace_ptr)); + CUTLASS_CHECK(gemm.run(stream)); + return cudaSuccess; +} + +} // namespace gemm + +} // namespace flashinfer + +#endif // FLASHINFER_GEMM_GROUPWISE_SM100_CUH_ diff --git a/include/flashinfer/gemm/group_gemm.cuh b/include/flashinfer/gemm/group_gemm.cuh new file mode 100644 index 0000000000000000000000000000000000000000..07ee09317e69e251b781ad0b44659cf95ddba41f --- /dev/null +++ b/include/flashinfer/gemm/group_gemm.cuh @@ -0,0 +1,115 @@ +/* + * Copyright (c) 2024 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef FLASHINFER_GEMM_GROUP_GEMM_CUH_ +#define FLASHINFER_GEMM_GROUP_GEMM_CUH_ + +#include + +#include "../allocator.h" +#include "../cutlass_utils.cuh" + +namespace flashinfer { + +namespace group_gemm { + +#define DISPATCH_WEIGHT_LAYOUT(is_column_major, WEIGHT_LAYOUT, ...) \ + if (is_column_major) { \ + using WEIGHT_LAYOUT = cutlass::layout::ColumnMajor; \ + __VA_ARGS__ \ + } else { \ + using WEIGHT_LAYOUT = cutlass::layout::RowMajor; \ + __VA_ARGS__ \ + } + +#define DISPATCH_SMEM_CONFIG(smem_limit_per_sm, NUM_STAGES, ...) \ + if (smem_limit_per_sm >= 147968) { \ + constexpr uint32_t NUM_STAGES = 4; \ + __VA_ARGS__ \ + } else { \ + constexpr uint32_t NUM_STAGES = 2; \ + __VA_ARGS__ \ + } + +template +cudaError_t CutlassSegmentGEMMRun(void* workspace_buffer, size_t workspace_buffer_size_in_bytes, + void* all_problems, int64_t batch_size, void* x, void* w, void* y, + void* x_ld, void* w_ld, void* y_ld, bool weight_column_major, + cudaStream_t stream) { + using cutlass::epilogue::thread::LinearCombination; + using cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle; + int device; + int smem_limit_per_sm; + cudaGetDevice(&device); + cudaDeviceGetAttribute(&smem_limit_per_sm, cudaDevAttrMaxSharedMemoryPerMultiprocessor, device); + + DISPATCH_WEIGHT_LAYOUT(weight_column_major, WEIGHT_LAYOUT, { + DISPATCH_SMEM_CONFIG(smem_limit_per_sm, NUM_STAGES, { + using GemmKernel = typename cutlass::gemm::kernel::DefaultGemmGrouped< + DType, // Element A + cutlass::layout::RowMajor, // Layout A + cutlass::ComplexTransform::kNone, // + 8, // Granularity A + DType, // Element B + WEIGHT_LAYOUT, // Layout B + cutlass::ComplexTransform::kNone, // + 8, // Granularity B + DType, // Element C&D + cutlass::layout::RowMajor, // Layout C&D + float, // Element Accumulator + cutlass::arch::OpClassTensorOp, // Operator Class Tag + cutlass::arch::Sm80, // Architecture + cutlass::gemm::GemmShape<128, 128, 32>, // Thread Block Shape + cutlass::gemm::GemmShape<64, 64, 32>, // Warp Shape + cutlass::gemm::GemmShape<16, 8, 16>, // Instruction Shape + cutlass::epilogue::thread::LinearCombination, // Epilogue + cutlass::gemm::threadblock::GemmBatchedIdentityThreadblockSwizzle, // Swizzling Operator + NUM_STAGES // Stages + >::GemmKernel; + + using EpilogueOutputOp = typename GemmKernel::Epilogue::OutputOp; + typename EpilogueOutputOp::Params epilogue_op(1.0, 1.0); + using GemmGrouped = cutlass::gemm::device::GemmGrouped; + typename GemmGrouped::Arguments args( + reinterpret_cast(all_problems), (int)batch_size, + /*threadblock_count=*/4, epilogue_op, static_cast(x), static_cast(w), + static_cast(y), static_cast(y), reinterpret_cast(x_ld), + reinterpret_cast(w_ld), reinterpret_cast(y_ld), + reinterpret_cast(y_ld)); + + GemmGrouped gemm; + auto status = gemm.initialize(args, nullptr, stream); + if (status != cutlass::Status::kSuccess) { + std::ostringstream err_msg; + err_msg << "cutlass group_gemm.initialize failed: " << cutlassGetStatusString(status); + FLASHINFER_ERROR(err_msg.str()); + } + status = gemm.run(stream); + if (status != cutlass::Status::kSuccess) { + std::ostringstream err_msg; + err_msg << "cutlass group_gemm.run failed: " << cutlassGetStatusString(status); + FLASHINFER_ERROR(err_msg.str()); + } + }); + }); + + return cudaSuccess; +} + +} // namespace group_gemm + +} // namespace flashinfer + +#endif // FLASHINFER_GEMM_GROUP_GEMM_CUH_ diff --git a/include/flashinfer/gemm/group_gemm_fp8_groupwise_sm100.cuh b/include/flashinfer/gemm/group_gemm_fp8_groupwise_sm100.cuh new file mode 100644 index 0000000000000000000000000000000000000000..03ece19e5f0a2295beb785bc846b998d16098f9f --- /dev/null +++ b/include/flashinfer/gemm/group_gemm_fp8_groupwise_sm100.cuh @@ -0,0 +1,255 @@ +/* + * Copyright (c) 2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef FLASHINFER_GROUP_GEMM_FP8_GROUPWISE_SM100_CUH_ +#define FLASHINFER_GROUP_GEMM_FP8_GROUPWISE_SM100_CUH_ + +#include +#include + +#include "../allocator.h" +#include "../cutlass_utils.cuh" +#include "../utils.cuh" + +namespace flashinfer { + +namespace group_gemm { + +using namespace cute; + +template +__global__ void compute_sm100_cutlass_group_gemm_args( + DTypeIn* A, DTypeIn* B, DTypeSF* SFA, DTypeSF* SFB, DTypeOut* D, int* m_indptr, int max_m, + int n, int k, int num_groups, int scale_granularity_m, int scale_granularity_n, + int scale_granularity_k, ProblemShape* problem_sizes, const DTypeIn** A_ptr, + const DTypeIn** B_ptr, const DTypeSF** SFA_ptr, const DTypeSF** SFB_ptr, DTypeOut** D_ptr, + StrideA* stride_A, StrideB* stride_B, StrideD* stride_D, LayoutSFA* layout_SFA, + LayoutSFB* layout_SFB) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= num_groups) { + return; + } + int sf_n = n / scale_granularity_n; + int sf_k = k / scale_granularity_k; +#if (__CUDACC_VER_MAJOR__ >= 12 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.wait;"); + asm volatile("griddepcontrol.launch_dependents;"); +#endif + int m_offset = m_indptr[i]; + int m_offset_next = m_indptr[i + 1]; + int m = m_offset_next - m_offset; + int sf_m_offset = m_offset / scale_granularity_m; + problem_sizes[i] = ProblemShape(m, n, k); + stride_A[i] = cutlass::make_cute_packed_stride(StrideA{}, {m, k, 1}); + stride_B[i] = cutlass::make_cute_packed_stride(StrideB{}, {n, k, 1}); + stride_D[i] = cutlass::make_cute_packed_stride(StrideD{}, {m, n, 1}); + A_ptr[i] = A + int64_t(m_offset) * int64_t(k); + B_ptr[i] = B + int64_t(i) * int64_t(n) * int64_t(k); + D_ptr[i] = D + int64_t(m_offset) * int64_t(n); + if constexpr (ScaleMajorK) { + layout_SFA[i] = ScaleConfig::tile_atom_to_shape_SFA(make_shape(m, n, k, 1)); + SFA_ptr[i] = SFA + int64_t(sf_m_offset) * int64_t(sf_k); + } else { + layout_SFA[i] = ScaleConfig::tile_atom_to_shape_SFA(make_shape(max_m, n, k, 1)); + SFA_ptr[i] = SFA + int64_t(sf_m_offset); + } + layout_SFB[i] = ScaleConfig::tile_atom_to_shape_SFB(make_shape(m, n, k, 1)); + SFB_ptr[i] = SFB + int64_t(i) * int64_t(sf_n) * int64_t(sf_k); +} + +template +cudaError_t CutlassFP8GroupwiseScaledGroupGEMMSM100( + void* int_buffer, size_t int_buffer_size_in_bytes, void* float_buffer, + size_t float_buffer_size_in_bytes, DTypeIn* A, DTypeIn* B, float* SFA, float* SFB, DTypeOut* D, + int* m_indptr, int max_m, int n, int k, int num_groups, cudaStream_t stream) { + using ProblemShape = cutlass::gemm::GroupProblemShape>; // per group + + using ElementA = DTypeIn; // Element type for A matrix operand + using LayoutA = cutlass::layout::RowMajor; // Layout type for A matrix operand + constexpr int AlignmentA = + 128 / cutlass::sizeof_bits::value; // Memory access granularity/alignment of A + // matrix in units of elements (up to 16 bytes) + + using ElementB = DTypeIn; // Element type for B matrix operand + using LayoutB = cutlass::layout::ColumnMajor; // Layout type for B matrix operand + constexpr int AlignmentB = + 128 / cutlass::sizeof_bits::value; // Memory access granularity/alignment of B + // matrix in units of elements (up to 16 bytes) + + using ElementD = DTypeOut; // Element type for C and D matrix operands + using LayoutD = cutlass::layout::RowMajor; // Layout type for C and D matrix operands + constexpr int AlignmentD = + 128 / cutlass::sizeof_bits::value; // Memory access granularity/alignment of D + // matrix in units of elements (up to 16 bytes) + + using ElementC = void; + using LayoutC = void; + constexpr int AlignmentC = 0; + + using ElementAccumulator = float; + using ElementCompute = float; + + using MmaTileShape_MNK = Shape, _128, _128>; + using ClusterShape_MNK = Shape, _1, _1>; + + using ScaleConfig = std::conditional_t< + ScaleMajorK, + cutlass::detail::Sm100BlockwiseScaleConfig, + cutlass::detail::Sm100BlockwiseScaleConfig>; + + using LayoutSFA = decltype(ScaleConfig::deduce_layoutSFA()); + using LayoutSFB = decltype(ScaleConfig::deduce_layoutSFB()); + + using EpilogueSchedule = + std::conditional_t; + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, MmaTileShape_MNK, ClusterShape_MNK, + cutlass::epilogue::collective::EpilogueTileAuto, ElementAccumulator, ElementCompute, ElementC, + LayoutC*, AlignmentC, ElementD, LayoutD*, AlignmentD, EpilogueSchedule>::CollectiveOp; + + using MainloopSchedule = + std::conditional_t; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, ElementA, + cute::tuple, AlignmentA, ElementB, cute::tuple, + AlignmentB, ElementAccumulator, MmaTileShape_MNK, ClusterShape_MNK, + cutlass::gemm::collective::StageCountAutoCarveout( + sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule>::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + + using StrideA = typename Gemm::GemmKernel::InternalStrideA; + using StrideB = typename Gemm::GemmKernel::InternalStrideB; + using StrideD = typename Gemm::GemmKernel::InternalStrideD; + + static_assert( + cute::is_same_v); + static_assert( + cute::is_same_v); + + AlignedAllocator allocator(int_buffer, int_buffer_size_in_bytes); + + auto problem_sizes = allocator.aligned_alloc( + num_groups * sizeof(typename ProblemShape::UnderlyingProblemShape), 16, + "sm100_groupwise_group_gemm_problem_sizes"); + auto A_ptr = allocator.aligned_alloc( + num_groups * sizeof(const typename Gemm::ElementA*), 16, "sm100_groupwise_group_gemm_A_ptr"); + auto B_ptr = allocator.aligned_alloc( + num_groups * sizeof(const typename Gemm::ElementB*), 16, "sm100_groupwise_group_gemm_B_ptr"); + auto D_ptr = allocator.aligned_alloc( + num_groups * sizeof(typename Gemm::EpilogueOutputOp::ElementOutput*), 16, + "sm100_groupwise_group_gemm_D_ptr"); + auto SFA_ptr = allocator.aligned_alloc( + num_groups * sizeof(const ElementAccumulator*), 16, "sm100_groupwise_group_gemm_SFA_ptr"); + auto SFB_ptr = allocator.aligned_alloc( + num_groups * sizeof(const ElementAccumulator*), 16, "sm100_groupwise_group_gemm_SFB_ptr"); + + auto stride_A = allocator.aligned_alloc(num_groups * sizeof(StrideA), 16, + "sm100_groupwise_group_gemm_stride_A"); + auto stride_B = allocator.aligned_alloc(num_groups * sizeof(StrideB), 16, + "sm100_groupwise_group_gemm_stride_B"); + auto stride_D = allocator.aligned_alloc(num_groups * sizeof(StrideD), 16, + "sm100_groupwise_group_gemm_stride_D"); + auto layout_SFA = allocator.aligned_alloc(num_groups * sizeof(LayoutSFA), 16, + "sm100_groupwise_group_gemm_layout_SFA"); + auto layout_SFB = allocator.aligned_alloc(num_groups * sizeof(LayoutSFB), 16, + "sm100_groupwise_group_gemm_layout_SFB"); + + int num_threads = std::min(num_groups, 1024); + int num_blocks = (num_groups + num_threads - 1) / num_threads; + cudaLaunchConfig_t config; + config.gridDim = num_blocks; + config.blockDim = num_threads; + config.dynamicSmemBytes = 0; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = true; + config.numAttrs = 1; + config.attrs = attrs; + + auto prepare_args_kernel = + compute_sm100_cutlass_group_gemm_args; + + FLASHINFER_CUDA_CALL(cudaLaunchKernelEx( + &config, prepare_args_kernel, A, B, SFA, SFB, D, m_indptr, max_m, n, k, num_groups, + ScaleGranularityM, ScaleGranularityN, ScaleGranularityK, problem_sizes, A_ptr, B_ptr, SFA_ptr, + SFB_ptr, D_ptr, stride_A, stride_B, stride_D, layout_SFA, layout_SFB)); + + thread_local int const sm_count = + cutlass::KernelHardwareInfo::query_device_multiprocessor_count(); + cutlass::KernelHardwareInfo hw_info; + hw_info.device_id = 0; + hw_info.sm_count = sm_count; + + typename Gemm::Arguments arguments{cutlass::gemm::GemmUniversalMode::kGrouped, + {num_groups, problem_sizes, /*problem_sizes_host=*/nullptr}, + { + A_ptr, + stride_A, + B_ptr, + stride_B, + SFA_ptr, + layout_SFA, + SFB_ptr, + layout_SFB, + }, + { + {}, // epilogue.thread + nullptr, // C_ptr + nullptr, // stride_C + D_ptr, + stride_D, + }, + hw_info}; + auto& fusion_args = arguments.epilogue.thread; + fusion_args.alpha = 1.0f; + fusion_args.beta = 0.0f; + + Gemm gemm; + + size_t workspace_size = Gemm::get_workspace_size(arguments); + AlignedAllocator float_allocator(float_buffer, float_buffer_size_in_bytes); + auto workspace_ptr = float_allocator.aligned_alloc( + workspace_size, 16, "sm100_groupwise_group_gemm_float_workspace"); + + CUTLASS_CHECK(gemm.can_implement(arguments)); + CUTLASS_CHECK(gemm.initialize(arguments, workspace_ptr)); + CUTLASS_CHECK(gemm.run(stream, /*cuda_adapter=*/nullptr, /*launch_with_pdl=*/true)); + return cudaSuccess; +} + +} // namespace group_gemm + +} // namespace flashinfer + +#endif // FLASHINFER_GROUP_GEMM_FP8_GROUPWISE_SM100_CUH_ diff --git a/include/flashinfer/gemm/group_gemm_lora.cuh b/include/flashinfer/gemm/group_gemm_lora.cuh new file mode 100644 index 0000000000000000000000000000000000000000..517419da5d502050036047c3ecafedaa99535b6d --- /dev/null +++ b/include/flashinfer/gemm/group_gemm_lora.cuh @@ -0,0 +1,29 @@ +/* + * Copyright (c) 2024 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef FLASHINFER_GROUP_GEMM_LORA_CUH_ +#define FLASHINFER_GROUP_GEMM_LORA_CUH_ + +namespace flashinfer { + +namespace group_gemm { + +// TODO(Zihao): port punica's sgmv kernel + +} // namespace group_gemm + +} // namespace flashinfer + +#endif // FLASHINFER_GROUP_GEMM_LORA_CUH_ diff --git a/include/flashinfer/gemm/group_gemm_mxfp4_groupwise_sm100.cuh b/include/flashinfer/gemm/group_gemm_mxfp4_groupwise_sm100.cuh new file mode 100644 index 0000000000000000000000000000000000000000..8856a218a1b92ed1a911686d50687d3727d148f7 --- /dev/null +++ b/include/flashinfer/gemm/group_gemm_mxfp4_groupwise_sm100.cuh @@ -0,0 +1,291 @@ +/* + * Copyright (c) 2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef FLASHINFER_GROUP_GEMM_MXFP4_GROUPWISE_SM100_CUH_ +#define FLASHINFER_GROUP_GEMM_MXFP4_GROUPWISE_SM100_CUH_ + +#include +#include + +#include "../allocator.h" +#include "../cutlass_utils.cuh" +#include "../utils.cuh" + +namespace flashinfer { + +namespace group_gemm { + +using namespace cute; + +// Function to safely offset an pointer that may contain sub-byte types (FP4/INT4) +template +__host__ __device__ __forceinline__ constexpr T* safe_inc_ptr(T* ptr, size_t offset) { + constexpr int adjustment = (sizeof_bits::value < 8) ? (8 / sizeof_bits::value) : 1; + assert(offset % adjustment == 0 && "Attempt to offset index to sub-byte"); + return ptr + offset / adjustment; +} + +template +__global__ void compute_sm100_cutlass_group_gemm_args( + ElementA* A, ElementB* B, ElementSFA* SFA, ElementSFB* SFB, ElementD* D, int* m_indptr, int n, + int k, int num_groups, ProblemShape* problem_sizes, const ElementA** A_ptr, + const ElementB** B_ptr, const ElementSFA** SFA_ptr, const ElementSFB** SFB_ptr, + ElementD** D_ptr, StrideA* stride_A, StrideB* stride_B, StrideD* stride_D, + LayoutSFA* layout_SFA, LayoutSFB* layout_SFB) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= num_groups) { + return; + } + constexpr int alignment_swizzled_mn = 128; + constexpr int alignment_swizzled_k = ScaleGranularity * 4; + int sf_n = (n + alignment_swizzled_mn - 1) / alignment_swizzled_mn * alignment_swizzled_mn; + int swizzled_k = (k + alignment_swizzled_k - 1) / alignment_swizzled_k * alignment_swizzled_k; + int sf_k = swizzled_k / ScaleGranularity; +#if (__CUDACC_VER_MAJOR__ >= 12 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.wait;"); + asm volatile("griddepcontrol.launch_dependents;"); +#endif + int m_offset = m_indptr[i]; + int m_offset_next = m_indptr[i + 1]; + int m = m_offset_next - m_offset; + // This formulation ensures that sf_m_offset_next - sf_m_offset >= m_offset_next - m_offset + int sf_m_offset = + (m_offset + i * (alignment_swizzled_mn - 1)) / alignment_swizzled_mn * alignment_swizzled_mn; + if constexpr (SwapAB) { + problem_sizes[i] = ProblemShape(n, m, k); + stride_A[i] = cutlass::make_cute_packed_stride(StrideA{}, {n, k, 1}); + stride_B[i] = cutlass::make_cute_packed_stride(StrideB{}, {m, k, 1}); + stride_D[i] = cutlass::make_cute_packed_stride(StrideD{}, {n, m, 1}); + A_ptr[i] = safe_inc_ptr(A, int64_t(i) * int64_t(n) * int64_t(k)); + B_ptr[i] = safe_inc_ptr(B, int64_t(m_offset) * int64_t(k)); + D_ptr[i] = safe_inc_ptr(D, int64_t(n) * int64_t(m_offset)); + layout_SFA[i] = ScaleConfig::tile_atom_to_shape_SFA(make_shape(sf_n, m, swizzled_k, 1)); + SFA_ptr[i] = safe_inc_ptr(SFA, int64_t(i) * int64_t(sf_n) * int64_t(sf_k)); + layout_SFB[i] = ScaleConfig::tile_atom_to_shape_SFB(make_shape(sf_n, m, swizzled_k, 1)); + SFB_ptr[i] = safe_inc_ptr(SFB, int64_t(sf_m_offset) * int64_t(sf_k)); + } else { + problem_sizes[i] = ProblemShape(m, n, k); + stride_A[i] = cutlass::make_cute_packed_stride(StrideA{}, {m, k, 1}); + stride_B[i] = cutlass::make_cute_packed_stride(StrideB{}, {n, k, 1}); + stride_D[i] = cutlass::make_cute_packed_stride(StrideD{}, {m, n, 1}); + A_ptr[i] = safe_inc_ptr(A, int64_t(m_offset) * int64_t(k)); + B_ptr[i] = safe_inc_ptr(B, int64_t(i) * int64_t(n) * int64_t(k)); + D_ptr[i] = safe_inc_ptr(D, int64_t(m_offset) * int64_t(n)); + layout_SFA[i] = ScaleConfig::tile_atom_to_shape_SFA(make_shape(m, sf_n, swizzled_k, 1)); + SFA_ptr[i] = safe_inc_ptr(SFA, int64_t(sf_m_offset) * int64_t(sf_k)); + layout_SFB[i] = ScaleConfig::tile_atom_to_shape_SFB(make_shape(m, sf_n, swizzled_k, 1)); + SFB_ptr[i] = safe_inc_ptr(SFB, int64_t(i) * int64_t(sf_n) * int64_t(sf_k)); + } +} + +template +cudaError_t CutlassMXFP4GroupwiseScaledGroupGEMMSM100( + void* int_buffer, size_t int_buffer_size_in_bytes, void* float_buffer, + size_t float_buffer_size_in_bytes, DTypeInA* A, DTypeInB* B, DTypeSFA* SFA, DTypeSFB* SFB, + DTypeOut* D, int* m_indptr, int n, int k, int num_groups, cudaStream_t stream) { + using ProblemShape = cutlass::gemm::GroupProblemShape>; // per group + + using ElementA = + std::conditional_t; // Element type for A matrix operand + using ElementSFA = + std::conditional_t; // Element type for SFA matrix operand + using LayoutA = cutlass::layout::RowMajor; // Layout type for A matrix operand + constexpr int AlignmentA = + cutlass::is_subbyte::value + ? 128 + : 128 / cutlass::sizeof_bits::value; // Memory access granularity/alignment of + // A matrix in units of elements + + using ElementB = + std::conditional_t; // Element type for B matrix operand + using ElementSFB = + std::conditional_t; // Element type for SFB matrix operand + using LayoutB = cutlass::layout::ColumnMajor; // Layout type for B matrix operand + constexpr int AlignmentB = + cutlass::is_subbyte::value + ? 128 + : 128 / cutlass::sizeof_bits::value; // Memory access granularity/alignment of + // B matrix in units of elements + + using ElementD = DTypeOut; // Element type for D matrix operands + using LayoutD = + std::conditional_t; // Layout type for D matrix operands + constexpr int AlignmentD = + 128 / cutlass::sizeof_bits::value; // Memory access granularity/alignment of D + // matrix in units of elements + + using ElementC = void; + using LayoutC = void; + constexpr int AlignmentC = 0; + + FLASHINFER_CHECK(k % std::max(AlignmentA, AlignmentB) == 0, "k must be divisible by %d", + max(AlignmentA, AlignmentB)); + FLASHINFER_CHECK(n % AlignmentD == 0, "n must be divisible by %d", AlignmentD); + + using ElementAccumulator = float; + using ElementCompute = float; + + using MmaTileShape_MNK = Shape, cute::Int, cute::Int>; + using ClusterShape_MNK = Shape, _1, _1>; + + using EpilogueSchedule = + std::conditional_t; + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp, MmaTileShape_MNK, ClusterShape_MNK, + cutlass::epilogue::collective::EpilogueTileAuto, ElementAccumulator, ElementCompute, ElementC, + LayoutC*, AlignmentC, ElementD, LayoutD*, AlignmentD, EpilogueSchedule>::CollectiveOp; + + using MainloopSchedule = + std::conditional_t; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Sm100, cutlass::arch::OpClassBlockScaledTensorOp, + cute::tuple, LayoutA*, AlignmentA, cute::tuple, + LayoutB*, AlignmentB, ElementAccumulator, MmaTileShape_MNK, ClusterShape_MNK, + cutlass::gemm::collective::StageCountAutoCarveout( + sizeof(typename CollectiveEpilogue::SharedStorage))>, + MainloopSchedule>::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal; + + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + + using StrideA = typename Gemm::GemmKernel::InternalStrideA; + using StrideB = typename Gemm::GemmKernel::InternalStrideB; + using StrideD = typename Gemm::GemmKernel::InternalStrideD; + + using ScaleConfig = typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig; + + using LayoutSFA = typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFA; + using LayoutSFB = typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFB; + + constexpr int ScaleGranularity = Gemm::GemmKernel::CollectiveMainloop::TiledMma::SFVecSize; + static_assert(ScaleGranularity == 32); + + AlignedAllocator allocator(int_buffer, int_buffer_size_in_bytes); + + auto problem_sizes = allocator.aligned_alloc( + num_groups * sizeof(typename ProblemShape::UnderlyingProblemShape), 16, + "sm100_groupwise_group_gemm_problem_sizes"); + auto A_ptr = allocator.aligned_alloc( + num_groups * sizeof(const typename Gemm::ElementA*), 16, "sm100_groupwise_group_gemm_A_ptr"); + auto B_ptr = allocator.aligned_alloc( + num_groups * sizeof(const typename Gemm::ElementB*), 16, "sm100_groupwise_group_gemm_B_ptr"); + auto D_ptr = allocator.aligned_alloc( + num_groups * sizeof(typename Gemm::EpilogueOutputOp::ElementOutput*), 16, + "sm100_groupwise_group_gemm_D_ptr"); + auto SFA_ptr = allocator.aligned_alloc( + num_groups * sizeof(const ElementSFA*), 16, "sm100_groupwise_group_gemm_SFA_ptr"); + auto SFB_ptr = allocator.aligned_alloc( + num_groups * sizeof(const ElementSFB*), 16, "sm100_groupwise_group_gemm_SFB_ptr"); + + auto stride_A = allocator.aligned_alloc(num_groups * sizeof(StrideA), 16, + "sm100_groupwise_group_gemm_stride_A"); + auto stride_B = allocator.aligned_alloc(num_groups * sizeof(StrideB), 16, + "sm100_groupwise_group_gemm_stride_B"); + auto stride_D = allocator.aligned_alloc(num_groups * sizeof(StrideD), 16, + "sm100_groupwise_group_gemm_stride_D"); + auto layout_SFA = allocator.aligned_alloc(num_groups * sizeof(LayoutSFA), 16, + "sm100_groupwise_group_gemm_layout_SFA"); + auto layout_SFB = allocator.aligned_alloc(num_groups * sizeof(LayoutSFB), 16, + "sm100_groupwise_group_gemm_layout_SFB"); + + int num_threads = std::min(num_groups, 1024); + int num_blocks = (num_groups + num_threads - 1) / num_threads; + cudaLaunchConfig_t config; + config.gridDim = num_blocks; + config.blockDim = num_threads; + config.dynamicSmemBytes = 0; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = true; + config.numAttrs = 1; + config.attrs = attrs; + + auto prepare_args_kernel = compute_sm100_cutlass_group_gemm_args< + SwapAB, ScaleGranularity, ScaleConfig, ElementA, ElementB, ElementSFA, ElementSFB, ElementD, + ProblemShape::UnderlyingProblemShape, StrideA, StrideB, StrideD, LayoutSFA, LayoutSFB>; + + if constexpr (SwapAB) { + FLASHINFER_CUDA_CALL(cudaLaunchKernelEx(&config, prepare_args_kernel, B, A, SFB, SFA, D, + m_indptr, n, k, num_groups, problem_sizes, A_ptr, B_ptr, + SFA_ptr, SFB_ptr, D_ptr, stride_A, stride_B, stride_D, + layout_SFA, layout_SFB)); + } else { + FLASHINFER_CUDA_CALL(cudaLaunchKernelEx(&config, prepare_args_kernel, A, B, SFA, SFB, D, + m_indptr, n, k, num_groups, problem_sizes, A_ptr, B_ptr, + SFA_ptr, SFB_ptr, D_ptr, stride_A, stride_B, stride_D, + layout_SFA, layout_SFB)); + } + + thread_local int const sm_count = + cutlass::KernelHardwareInfo::query_device_multiprocessor_count(); + cutlass::KernelHardwareInfo hw_info; + hw_info.device_id = 0; + hw_info.sm_count = sm_count; + + typename Gemm::Arguments arguments{cutlass::gemm::GemmUniversalMode::kGrouped, + {num_groups, problem_sizes, /*problem_sizes_host=*/nullptr}, + { + A_ptr, + stride_A, + B_ptr, + stride_B, + SFA_ptr, + layout_SFA, + SFB_ptr, + layout_SFB, + }, + { + {}, // epilogue.thread + nullptr, // C_ptr + nullptr, // stride_C + D_ptr, + stride_D, + }, + hw_info}; + auto& fusion_args = arguments.epilogue.thread; + fusion_args.alpha = 1.0f; + fusion_args.beta = 0.0f; + + Gemm gemm; + + size_t workspace_size = Gemm::get_workspace_size(arguments); + AlignedAllocator float_allocator(float_buffer, float_buffer_size_in_bytes); + auto workspace_ptr = float_allocator.aligned_alloc( + workspace_size, 16, "sm100_groupwise_group_gemm_float_workspace"); + + CUTLASS_CHECK(gemm.can_implement(arguments)); + CUTLASS_CHECK(gemm.initialize(arguments, workspace_ptr)); + CUTLASS_CHECK(gemm.run(stream, /*cuda_adapter=*/nullptr, /*launch_with_pdl=*/true)); + return cudaSuccess; +} + +} // namespace group_gemm + +} // namespace flashinfer + +#endif // FLASHINFER_GROUP_GEMM_MXFP4_GROUPWISE_SM100_CUH_ diff --git a/include/flashinfer/gemm/group_gemm_sm90.cuh b/include/flashinfer/gemm/group_gemm_sm90.cuh new file mode 100644 index 0000000000000000000000000000000000000000..a54e1fa0e71928717b56494204f2e8cfd6820b5d --- /dev/null +++ b/include/flashinfer/gemm/group_gemm_sm90.cuh @@ -0,0 +1,160 @@ +/* + * Copyright (c) 2024 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef FLASHINFER_GEMM_GROUP_GEMM_SM90_CUH_ +#define FLASHINFER_GEMM_GROUP_GEMM_SM90_CUH_ + +#include "../allocator.h" +#include "../cutlass_utils.cuh" +#include "../utils.cuh" + +namespace flashinfer { + +namespace group_gemm { + +using namespace cute; + +#define DISPATCH_WEIGHT_LAYOUT(is_column_major, WEIGHT_LAYOUT, ...) \ + if (is_column_major) { \ + using WEIGHT_LAYOUT = cutlass::layout::ColumnMajor; \ + __VA_ARGS__ \ + } else { \ + using WEIGHT_LAYOUT = cutlass::layout::RowMajor; \ + __VA_ARGS__ \ + } + +/** + * Panic wrapper for unwinding CUTLASS errors + */ +#define CUTLASS_CHECK(status) \ + { \ + cutlass::Status error = status; \ + if (error != cutlass::Status::kSuccess) { \ + std::cerr << "Got cutlass error: " << cutlassGetStatusString(error) << " at: " << __LINE__ \ + << std::endl; \ + exit(EXIT_FAILURE); \ + } \ + } + +template +cudaError_t CutlassSegmentGEMMSM90Run(void* float_buffer, size_t float_buffer_size_in_bytes, + void* int_buffer, size_t int_buffer_size_in_bytes, + void* all_problems, int64_t batch_size, void* x, void* w, + void* y, void* x_stride, void* w_stride, void* y_stride, + bool weight_column_major, cudaStream_t stream) { + auto compute_capacity = GetCudaComputeCapability(); + if (compute_capacity.first < 9) { + std::cerr << "CutlassSegmentGEMMSM90Run requires compute capability of at least 9.0" + << std::endl; + return cudaErrorNotSupported; + } + + using ProblemShape = cutlass::gemm::GroupProblemShape>; + using ElementA = DTypeIn; + using ElementB = DTypeIn; + using ElementC = DTypeOut; + + DISPATCH_WEIGHT_LAYOUT(weight_column_major, WEIGHT_LAYOUT, { + if constexpr (std::is_same_v && + sizeof(DTypeIn) == 1) { + std::ostringstream err_msg; + err_msg << "Row-major layout is not supported for fp8 data type"; + FLASHINFER_ERROR(err_msg.str()); + } else { + using LayoutA = cutlass::layout::RowMajor; + constexpr int AlignmentA = 128 / cutlass::sizeof_bits::value; + + using LayoutB = WEIGHT_LAYOUT; + constexpr int AlignmentB = 128 / cutlass::sizeof_bits::value; + + using LayoutC = cutlass::layout::RowMajor; + constexpr int AlignmentC = 128 / cutlass::sizeof_bits::value; + + constexpr bool is_fp8 = sizeof(DTypeIn) == 1; + + using ElementAccumulator = float; + using ArchTag = cutlass::arch::Sm90; + using OperatorClass = cutlass::arch::OpClassTensorOp; + using TileShape = + typename std::conditional, Shape<_128, _128, _128>>::type; + using ClusterShape = + typename std::conditional, Shape<_2, _1, _1>>::type; + using StageCountType = cutlass::gemm::collective::StageCountAuto; + using KernelSchedule = typename std::conditional< + is_fp8, cutlass::gemm::KernelPtrArrayTmaWarpSpecializedCooperativeFP8FastAccum, + cutlass::gemm::KernelPtrArrayTmaWarpSpecializedCooperative>::type; + using EpilogueSchedule = cutlass::epilogue::PtrArrayNoSmemWarpSpecialized; + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Sm90, cutlass::arch::OpClassTensorOp, TileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, ElementAccumulator, ElementAccumulator, + ElementC, LayoutC*, AlignmentC, ElementC, LayoutC*, AlignmentC, + EpilogueSchedule>::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + ArchTag, OperatorClass, ElementA, LayoutA*, AlignmentA, ElementB, LayoutB*, AlignmentB, + ElementAccumulator, TileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout( + sizeof(typename CollectiveEpilogue::SharedStorage))>, + KernelSchedule>::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + + using StrideA = typename Gemm::GemmKernel::InternalStrideA; + using StrideB = typename Gemm::GemmKernel::InternalStrideB; + using StrideC = typename Gemm::GemmKernel::InternalStrideC; + using StrideD = typename Gemm::GemmKernel::InternalStrideD; + + cutlass::KernelHardwareInfo hw_info; + cudaGetDevice(&hw_info.device_id); + hw_info.sm_count = + cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id); + + typename Gemm::EpilogueOutputOp::Params params; + params = + typename Gemm::EpilogueOutputOp::Params(ElementAccumulator(1.f), ElementAccumulator(0.f)); + + typename Gemm::Arguments arguments{ + cutlass::gemm::GemmUniversalMode::kGrouped, + {int(batch_size), reinterpret_cast(all_problems), + nullptr}, + {static_cast(x), reinterpret_cast(x_stride), + static_cast(w), reinterpret_cast(w_stride)}, + {params, static_cast(y), reinterpret_cast(y_stride), + static_cast(y), reinterpret_cast(y_stride)}, + hw_info}; + + Gemm gemm; + + size_t workspace_size = Gemm::get_workspace_size(arguments); + AlignedAllocator float_allocator(float_buffer, float_buffer_size_in_bytes); + auto workspace_ptr = float_allocator.aligned_alloc(workspace_size, 64, + "sm90_group_gemm_float_workspace"); + + CUTLASS_CHECK(gemm.can_implement(arguments)); + CUTLASS_CHECK(gemm.initialize(arguments, workspace_ptr)); + CUTLASS_CHECK(gemm.run(stream)); + } + }); + + return cudaSuccess; +} + +} // namespace group_gemm +} // namespace flashinfer + +#endif // FLASHINFER_GEMM_GROUP_GEMM_SM90_CUH_ diff --git a/include/flashinfer/gemm/group_gemv.cuh b/include/flashinfer/gemm/group_gemv.cuh new file mode 100644 index 0000000000000000000000000000000000000000..4b439355e4a5c5a842440ede2011798f42b807aa --- /dev/null +++ b/include/flashinfer/gemm/group_gemv.cuh @@ -0,0 +1,29 @@ +/* + * Copyright (c) 2024 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef FLASHINFER_GROUP_GEMV_CUH_ +#define FLASHINFER_GROUP_GEMV_CUH_ + +namespace flashinfer { + +namespace group_gemm { + +// TODO(Zihao): port punica's bgmv kernel + +} // namespace group_gemm + +} // namespace flashinfer + +#endif // FLASHINFER_GROUP_GEMV_CUH_ diff --git a/include/flashinfer/layout.cuh b/include/flashinfer/layout.cuh new file mode 100644 index 0000000000000000000000000000000000000000..aeaf26b580c753293773b0c73d45a3badacee720 --- /dev/null +++ b/include/flashinfer/layout.cuh @@ -0,0 +1,128 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef FLASHINFER_LAYOUT_CUH_ +#define FLASHINFER_LAYOUT_CUH_ + +#include +#include +#include + +namespace flashinfer { + +/*! + * \brief The Layout of QKV matrices + */ +enum class QKVLayout { + // [seq_len, num_heads, head_dim] + kNHD = 0U, + // [num_heads, seq_len, head_dim] + kHND = 1U, +}; + +__host__ __device__ __forceinline__ size_t get_elem_offset_impl(size_t elem_idx, size_t head_idx, + size_t feat_idx, size_t stride_n, + size_t stride_h) { + return elem_idx * stride_n + head_idx * stride_h + feat_idx; +} + +__host__ __forceinline__ auto get_qkv_strides(QKVLayout kv_layout, uint32_t kv_len, + uint32_t num_qo_heads, uint32_t num_kv_heads, + uint32_t head_dim) { + const uint32_t q_stride_n = num_qo_heads * head_dim, q_stride_h = head_dim, + kv_stride_n = (kv_layout == QKVLayout::kNHD) ? num_kv_heads * head_dim : head_dim, + kv_stride_h = (kv_layout == QKVLayout::kNHD) ? head_dim : kv_len * head_dim; + return std::make_tuple(q_stride_n, q_stride_h, kv_stride_n, kv_stride_h); +} + +struct tensor_info_t { + uint32_t qo_len; + uint32_t kv_len; + uint32_t num_qo_heads; + uint32_t num_kv_heads; + uint32_t q_stride_n; + uint32_t q_stride_h; + uint32_t kv_stride_n; + uint32_t kv_stride_h; + uint32_t head_dim; + __host__ __device__ __forceinline__ tensor_info_t(uint32_t qo_len, uint32_t kv_len, + uint32_t num_qo_heads, uint32_t num_kv_heads, + uint32_t q_stride_n, uint32_t q_stride_h, + uint32_t kv_stride_n, uint32_t kv_stride_h, + uint32_t head_dim) + : qo_len(qo_len), + kv_len(kv_len), + num_qo_heads(num_qo_heads), + num_kv_heads(num_kv_heads), + q_stride_n(q_stride_n), + q_stride_h(q_stride_h), + kv_stride_n(kv_stride_n), + kv_stride_h(kv_stride_h), + head_dim(head_dim) {} + + __host__ __device__ __forceinline__ tensor_info_t(uint32_t qo_len, uint32_t kv_len, + uint32_t num_qo_heads, uint32_t num_kv_heads, + QKVLayout kv_layout, uint32_t head_dim) + : qo_len(qo_len), + kv_len(kv_len), + num_qo_heads(num_qo_heads), + num_kv_heads(num_kv_heads), + head_dim(head_dim) { + q_stride_n = num_qo_heads * head_dim; + q_stride_h = head_dim; + kv_stride_n = (kv_layout == QKVLayout::kNHD) ? num_kv_heads * head_dim : head_dim; + kv_stride_h = (kv_layout == QKVLayout::kNHD) ? head_dim : kv_len * head_dim; + } + + __host__ __device__ __forceinline__ size_t get_q_elem_offset(uint32_t qo_idx, + uint32_t qo_head_idx, + uint32_t feat_idx) const { + return get_elem_offset_impl(qo_idx, qo_head_idx, feat_idx, q_stride_n, q_stride_h); + } + + __host__ __device__ __forceinline__ size_t get_o_elem_offset(uint32_t qo_idx, + uint32_t qo_head_idx, + uint32_t feat_idx) const { + return get_elem_offset_impl(qo_idx, qo_head_idx, feat_idx, num_qo_heads * head_dim, head_dim); + } + + __host__ __device__ __forceinline__ size_t get_kv_elem_offset(uint32_t kv_idx, + uint32_t kv_head_idx, + uint32_t feat_idx) const { + return get_elem_offset_impl(kv_idx, kv_head_idx, feat_idx, kv_stride_n, kv_stride_h); + } + + __host__ __device__ __forceinline__ uint32_t get_group_size() const { + return num_qo_heads / num_kv_heads; + } +}; + +/*! + * \brief Convert QKVLayout to string + * \param layout The QKVLayout to convert + */ +inline std::string QKVLayoutToString(const QKVLayout& layout) { + switch (layout) { + case QKVLayout::kNHD: + return "NHD"; + case QKVLayout::kHND: + return "HND"; + default: + return "Unknown"; + } +} + +} // namespace flashinfer +#endif // FLASHINFER_LAYOUT_CUH_ diff --git a/include/flashinfer/logging.h b/include/flashinfer/logging.h new file mode 100644 index 0000000000000000000000000000000000000000..2bcacda704ee8bc7e67383f096deda5f400f1d71 --- /dev/null +++ b/include/flashinfer/logging.h @@ -0,0 +1,45 @@ +/* + * Copyright (c) 2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef FLASHINFER_LOGGING_H_ +#define FLASHINFER_LOGGING_H_ + +#include +#include + +#define FLASHINFER_LOG_TRACE(...) spdlog::trace(__VA_ARGS__) +#define FLASHINFER_LOG_DEBUG(...) spdlog::debug(__VA_ARGS__) +#define FLASHINFER_LOG_INFO(...) spdlog::info(__VA_ARGS__) +#define FLASHINFER_LOG_WARN(...) spdlog::warn(__VA_ARGS__) +#define FLASHINFER_LOG_ERROR(...) spdlog::error(__VA_ARGS__) +#define FLASHINFER_LOG_CRITICAL(...) spdlog::critical(__VA_ARGS__) + +namespace flashinfer { + +namespace logging { + +inline void set_log_level(spdlog::level::level_enum lvl) { + auto fmt = "[%Y-%m-%d %H:%M:%S.%f] [%n] [%^%l%$] %v"; + auto console_sink = std::make_shared(); + console_sink->set_pattern(fmt); + console_sink->set_level(lvl); + spdlog::set_default_logger(std::make_shared("flashinfer", console_sink)); +} + +} // namespace logging + +} // namespace flashinfer + +#endif // FLASHINFER_LOGGING_H_ diff --git a/include/flashinfer/math.cuh b/include/flashinfer/math.cuh new file mode 100644 index 0000000000000000000000000000000000000000..27c6351e8ff36cec490fe81ef5463699050c7a0e --- /dev/null +++ b/include/flashinfer/math.cuh @@ -0,0 +1,156 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef FLASHINFER_MATH_CUH_ +#define FLASHINFER_MATH_CUH_ + +#include +#include + +#include + +namespace flashinfer { +namespace math { + +// log2(e) +constexpr float log2e = 1.44269504088896340736f; + +constexpr float loge2 = 0.693147180559945309417f; + +constexpr float inf = 5e4; + +__forceinline__ __device__ half2 uint32_as_half2(uint32_t x) { return *(half2*)&x; } + +__forceinline__ __device__ uint32_t half2_as_uint32(half2 x) { return *(uint32_t*)&x; } + +/*! + * \brief Wrapper of PTX ex2.approx instruction, which computes 2^x + * \param x input + */ +__forceinline__ __device__ float ptx_exp2(float x) { + float y; + asm volatile("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + +/*! + * \brief Wrapper of PTX lg2.approx instruction, which computes log2(x) + * \param x input + */ +__forceinline__ __device__ float ptx_log2(float x) { + float y; + asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + +/*! + * \brief Wrapper of PTX ex2.approx.f16x2 instruction, which computes 2^x + * \param x input + */ +__forceinline__ __device__ half2 ptx_exp2(half2 x) { + uint32_t y_u32; + uint32_t x_u32 = half2_as_uint32(x); + asm volatile("ex2.approx.f16x2 %0, %1;" : "=r"(y_u32) : "r"(x_u32)); + return uint32_as_half2(y_u32); +} + +/*! + * \brief Wrapper of PTX ex2.approx.f16 instruction, which computes 2^x + * \param x input + */ +__forceinline__ __device__ half ptx_exp2(half x) { + ushort y_u16; + asm volatile("ex2.approx.f16 %0, %1;" : "=h"(y_u16) : "h"(__half_as_ushort(x))); + return __ushort_as_half(y_u16); +} + +/*! + * \brief Wrapper of PTX rcp.approx instruction, which computes 1/x + * \param x input + */ +__forceinline__ __device__ float ptx_rcp(float x) { + float y; + asm volatile("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + +/*! + * \brief Wrapper of PTX shfl.sync.bfly instruction, which performs a butterfly shuffle + * between threads in a warp. + * \param x The value in the source lane + * \param lane_mask The mask to perform thread index xor with: y[i] <- x[i ^ delta] + */ +__forceinline__ __device__ float shfl_xor_sync(float x, int lane_mask) { + float y; + asm volatile("shfl.sync.bfly.b32 %0, %1, %2, 0x1f, 0xffffffff;" + : "=f"(y) + : "f"(x), "r"(lane_mask)); + return y; +} + +/*! + * \brief Wrapper of PTX shfl.sync.bfly instruction on half2, which performs a butterfly + * shuffle between threads in a warp. + * \param x The value in the source lane + * \param lane_mask The mask to perform thread index xor with: y[i] <- x[i ^ lane_mask] + */ +__forceinline__ __device__ half2 shfl_xor_sync(half2 x, int lane_mask) { + return __shfl_xor_sync(0xffffffff, x, lane_mask); +} + +/*! + * \brief Wrapper of PTX rsqrt approximation instruction, which computes 1/sqrt(x) + * \param x input + */ +__forceinline__ __device__ float rsqrt(float x) { + float y; + asm volatile("rsqrt.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + +/*! + * \brief Wrapper of PTX tanh.approx.f32 instruction, which computes tanh(x) + * \param x input + */ +__forceinline__ __device__ float tanh(float x) { + float y; + asm volatile("tanh.approx.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + +/*! + * \brief Wrapper of PTX tanh.approx.f16x2 instruction, which computes tanh(x) + * \param x input + */ +__forceinline__ __device__ half2 tanh(half2 x) { + uint32_t y_u32; + uint32_t x_u32 = half2_as_uint32(x); + asm volatile("tanh.approx.f16x2 %0, %1;" : "=r"(y_u32) : "r"(x_u32)); + return uint32_as_half2(y_u32); +} + +/*! + * \brief Wrapper of PTX tanh.approx.f16 instruction, which computes tanh(x) + * \param x input + */ +__forceinline__ __device__ half tanh(half x) { + ushort y_u16; + asm volatile("tanh.approx.f16 %0, %1;" : "=h"(y_u16) : "h"(__half_as_ushort(x))); + return __ushort_as_half(y_u16); +} + +} // namespace math +} // namespace flashinfer +#endif // FLASHINFER_MATH_CUH_ diff --git a/include/flashinfer/mma.cuh b/include/flashinfer/mma.cuh new file mode 100644 index 0000000000000000000000000000000000000000..75c7dc6be7a7e90e81591c902ae13ab1108c05cc --- /dev/null +++ b/include/flashinfer/mma.cuh @@ -0,0 +1,705 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef FLASHINFER_MMA_CUH_ +#define FLASHINFER_MMA_CUH_ + +#include +#include +#include +#include + +#include + +namespace flashinfer { + +namespace mma { + +#if (__CUDACC_VER_MAJOR__ * 10000 + __CUDACC_VER_MINOR__ * 100 >= 120400) +#if (!defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 890)) +#define FLASHINFER_MMA_F8F8F32_M16N8K32_ENABLED +#endif +#endif + +#if (__CUDACC_VER_MAJOR__ >= 11) +#if (!defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 900)) +#define FLASHINFER_STMATRIX_M8N8X4_ENABLED +#endif +#if (!defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 800)) +#define FLASHINFER_MMA_F16F16F32_M16N8K16_ENABLED +#define FLASHINFER_MMA_F16F16F16_M16N8K16_ENABLED +#endif +#if (!defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 750)) +#define FLASHINFER_MMA_F16F16F32_M16N8K8_ENABLED +#define FLASHINFER_MMA_F16F16F16_M16N8K8_ENABLED +#define FLASHINFER_LDMATRIX_M8N8X4_ENABLED +#endif +#endif + +#if defined(__CUDA_ARCH__) +#define FLASHINFER_RUNTIME_ASSERT(x) __brkpt() +#else +#define FLASHINFER_RUNTIME_ASSERT(x) assert(0 && x) +#endif + +enum class MMAMode { + kInit = 0U, + kInplaceUpdate = 1U, +}; + +/*! + * \brief Wrapper of PTX ldmatrix m8n8.x4 instruction, loads data from shared memory + * to fragment + * \tparam T data type of the fragment + * \param R pointer to the fragment + * \param smem_ptr pointer to the shared memory + */ +template +__device__ __forceinline__ void ldmatrix_m8n8x4(uint32_t* R, T* smem_ptr) { +#ifdef FLASHINFER_LDMATRIX_M8N8X4_ENABLED + uint32_t smem_int_ptr = static_cast(__cvta_generic_to_shared(smem_ptr)); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(R[0]), "=r"(R[1]), "=r"(R[2]), "=r"(R[3]) + : "r"(smem_int_ptr)); +#else + FLASHINFER_RUNTIME_ASSERT("Unsupported CUDA architecture for ldmatrix instruction"); +#endif +} + +/*! + * \brief Wrapper of PTX ldmatrix m8n8.x4 instruction, loads data from shared memory + * to fragment + * \tparam T data type of the fragment + * \param R pointer to the fragment + * \param smem_ptr pointer to the shared memory + */ +template +__device__ __forceinline__ void ldmatrix_m8n8x4_left_half(uint32_t* R, T* smem_ptr) { +#ifdef FLASHINFER_LDMATRIX_M8N8X4_ENABLED + uint32_t smem_int_ptr = static_cast(__cvta_generic_to_shared(smem_ptr)); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, _, %1, _}, [%2];\n" + : "=r"(R[0]), "=r"(R[1]) + : "r"(smem_int_ptr)); +#else + FLASHINFER_RUNTIME_ASSERT("Unsupported CUDA architecture for ldmatrix instruction"); +#endif +} + +/*! + * \brief Wrapper of PTX ldmatrix m8n8.x4 instruction, loads data from shared memory + * to fragment + * \tparam T data type of the fragment + * \param R pointer to the fragment + * \param smem_ptr pointer to the shared memory + */ +template +__device__ __forceinline__ void ldmatrix_m8n8x4_right_half(uint32_t* R, T* smem_ptr) { +#ifdef FLASHINFER_LDMATRIX_M8N8X4_ENABLED + uint32_t smem_int_ptr = static_cast(__cvta_generic_to_shared(smem_ptr)); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {_, %0, _, %1}, [%2];\n" + : "=r"(R[0]), "=r"(R[1]) + : "r"(smem_int_ptr)); +#else + FLASHINFER_RUNTIME_ASSERT("Unsupported CUDA architecture for ldmatrix instruction"); +#endif +} + +/*! + * \brief Wrapper of PTX ldmatrix m8n8.x4 transposed instruction, loads data from + * shared memory to fragment and transposes the fragment + * \tparam T data type of the fragment + * \param R pointer to the fragment + * \param smem_ptr pointer to the shared memory + */ +template +__device__ __forceinline__ void ldmatrix_m8n8x4_trans(uint32_t* R, T* smem_ptr) { +#ifdef FLASHINFER_LDMATRIX_M8N8X4_ENABLED + uint32_t smem_int_ptr = static_cast(__cvta_generic_to_shared(smem_ptr)); + asm volatile("ldmatrix.sync.aligned.trans.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(R[0]), "=r"(R[1]), "=r"(R[2]), "=r"(R[3]) + : "r"(smem_int_ptr)); +#else + FLASHINFER_RUNTIME_ASSERT("Unsupported CUDA architecture for ldmatrix instruction"); +#endif +} + +/*! + * \brief Wrapper of PTX ldmatrix m8n8.x4 transposed instruction, loads data from + * shared memory to fragment and transposes the fragment + * \tparam T data type of the fragment + * \param R pointer to the fragment + * \param smem_ptr pointer to the shared memory + */ +template +__device__ __forceinline__ void ldmatrix_m8n8x4_trans_left_half(uint32_t* R, T* smem_ptr) { +#ifdef FLASHINFER_LDMATRIX_M8N8X4_ENABLED + uint32_t smem_int_ptr = static_cast(__cvta_generic_to_shared(smem_ptr)); + asm volatile("ldmatrix.sync.aligned.trans.m8n8.x4.shared.b16 {%0, %1, _, _}, [%2];\n" + : "=r"(R[0]), "=r"(R[1]) + : "r"(smem_int_ptr)); +#else + FLASHINFER_RUNTIME_ASSERT("Unsupported CUDA architecture for ldmatrix instruction"); +#endif +} + +/*! + * \brief Wrapper of PTX ldmatrix m8n8.x4 transposed instruction, loads data from + * shared memory to fragment and transposes the fragment + * \tparam T data type of the fragment + * \param R pointer to the fragment + * \param smem_ptr pointer to the shared memory + */ +template +__device__ __forceinline__ void ldmatrix_m8n8x4_trans_right_half(uint32_t* R, T* smem_ptr) { +#ifdef FLASHINFER_LDMATRIX_M8N8X4_ENABLED + uint32_t smem_int_ptr = static_cast(__cvta_generic_to_shared(smem_ptr)); + asm volatile("ldmatrix.sync.aligned.trans.m8n8.x4.shared.b16 {_, _, %0, %1}, [%2];\n" + : "=r"(R[0]), "=r"(R[1]) + : "r"(smem_int_ptr)); +#else + FLASHINFER_RUNTIME_ASSERT("Unsupported CUDA architecture for ldmatrix instruction"); +#endif +} + +/*! + * \brief Wrapper of PTX stmatrix m8n8.x4 instruction, stores data from fragment + * to shared memory + * \tparam T data type of the fragment + * \param R pointer to the fragment + * \param smem_ptr pointer to the shared memory + */ +template +__device__ __forceinline__ void stmatrix_m8n8x4(uint32_t* R, T* smem_ptr) { +#ifdef FLASHINFER_STMATRIX_M8N8X4_ENABLED + uint32_t smem_int_ptr = static_cast(__cvta_generic_to_shared(smem_ptr)); + asm volatile("stmatrix.sync.aligned.m8n8.x4.shared.b16 [%0], {%1, %2, %3, %4};\n" + : + : "r"(smem_int_ptr), "r"(R[0]), "r"(R[1]), "r"(R[2]), "r"(R[3])); +#else + // Fallback implementation, slower than PTX instruction + const uint32_t tx = threadIdx.x; + uint4 word; +#pragma unroll + for (uint32_t reg_id = 0; reg_id < 4; ++reg_id) { + word.x = __shfl_sync(0xffffffff, R[reg_id], (tx % 8) * 4); + word.y = __shfl_sync(0xffffffff, R[reg_id], (tx % 8) * 4 + 1); + word.z = __shfl_sync(0xffffffff, R[reg_id], (tx % 8) * 4 + 2); + word.w = __shfl_sync(0xffffffff, R[reg_id], (tx % 8) * 4 + 3); + if (tx / 8 == reg_id) { + *(uint4*)smem_ptr = word; + } + } +#endif +} + +/*! + * \brief Wrapper of two mma m16n8k32 instructions for row major and column major f8 matrix + * multiplication, accumulated in f32. + * \tparam T data type of the fragment + * \tparam mma_mode whether we are initializing the accumulator or updating it + * \param C pointer to the accumulator + * \param A pointer to the fragment of matrix A + * \param B pointer to the fragment of matrix B + */ +template +__device__ __forceinline__ void mma_sync_m16n16k32_row_col_f8f8f32(float* C, uint32_t* A, + uint32_t* B) { + static_assert(sizeof(T) == 1, "DType must be 8bit floating data type"); +#if defined(FLASHINFER_MMA_F8F8F32_M16N8K32_ENABLED) + if constexpr (mma_mode == MMAMode::kInit) { + if constexpr (std::is_same_v) { + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(0.f), "f"(0.f), + "f"(0.f), "f"(0.f)); + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=f"(C[4]), "=f"(C[5]), "=f"(C[6]), "=f"(C[7]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "f"(0.f), "f"(0.f), + "f"(0.f), "f"(0.f)); + } else { // e5m2 + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.f32.e5m2.e5m2.f32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(0.f), "f"(0.f), + "f"(0.f), "f"(0.f)); + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.f32.e5m2.e5m2.f32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=f"(C[4]), "=f"(C[5]), "=f"(C[6]), "=f"(C[7]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "f"(0.f), "f"(0.f), + "f"(0.f), "f"(0.f)); + } + } else { + if constexpr (std::is_same_v) { + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(C[0]), "f"(C[1]), + "f"(C[2]), "f"(C[3])); + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=f"(C[4]), "=f"(C[5]), "=f"(C[6]), "=f"(C[7]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "f"(C[4]), "f"(C[5]), + "f"(C[6]), "f"(C[7])); + } else { // e5m2 + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.f32.e5m2.e5m2.f32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(C[0]), "f"(C[1]), + "f"(C[2]), "f"(C[3])); + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.f32.e5m2.e5m2.f32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=f"(C[4]), "=f"(C[5]), "=f"(C[6]), "=f"(C[7]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "f"(C[4]), "f"(C[5]), + "f"(C[6]), "f"(C[7])); + } + } +#else + FLASHINFER_RUNTIME_ASSERT( + "fp8 mma instruction is only available for sm89, PTX 8.4+ and CUDA 12.4+"); +#endif +} + +/*! + * \brief Wrapper of two mma m16n8k16 instructions for row major and column major f16 matrix + * multiplication, accumulated in f32. + * \tparam T data type of the fragment + * \tparam mma_mode whether we are initializing the accumulator or updating it + * \param C pointer to the accumulator + * \param A pointer to the fragment of matrix A + * \param B pointer to the fragment of matrix B + */ +template +__device__ __forceinline__ void mma_sync_m16n16k16_row_col_f16f16f32(float* C, uint32_t* A, + uint32_t* B) { +#if defined(FLASHINFER_MMA_F16F16F32_M16N8K16_ENABLED) + if constexpr (mma_mode == MMAMode::kInit) { + if constexpr (std::is_same_v) { + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(0.f), "f"(0.f), + "f"(0.f), "f"(0.f)); + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=f"(C[4]), "=f"(C[5]), "=f"(C[6]), "=f"(C[7]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "f"(0.f), "f"(0.f), + "f"(0.f), "f"(0.f)); + } else { + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(0.f), "f"(0.f), + "f"(0.f), "f"(0.f)); + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=f"(C[4]), "=f"(C[5]), "=f"(C[6]), "=f"(C[7]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "f"(0.f), "f"(0.f), + "f"(0.f), "f"(0.f)); + } + } else { + if constexpr (std::is_same_v) { + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(C[0]), "f"(C[1]), + "f"(C[2]), "f"(C[3])); + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=f"(C[4]), "=f"(C[5]), "=f"(C[6]), "=f"(C[7]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "f"(C[4]), "f"(C[5]), + "f"(C[6]), "f"(C[7])); + } else { + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(C[0]), "f"(C[1]), + "f"(C[2]), "f"(C[3])); + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=f"(C[4]), "=f"(C[5]), "=f"(C[6]), "=f"(C[7]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "f"(C[4]), "f"(C[5]), + "f"(C[6]), "f"(C[7])); + } + } +#elif defined(FLASHINFER_MMA_F16F16F32_M16N8K8_ENABLED) + if constexpr (std::is_same_v) { + if constexpr (mma_mode == MMAMode::kInit) { + asm volatile( + "mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 " + "{%0, %1, %2, %3}," + "{%4, %5}," + "{%6}," + "{%7, %8, %9, %10};\n" + : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3]) + : "r"(A[0]), "r"(A[1]), "r"(B[0]), "f"(0.f), "f"(0.f), "f"(0.f), "f"(0.f)); + asm volatile( + "mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 " + "{%0, %1, %2, %3}," + "{%4, %5}," + "{%6}," + "{%7, %8, %9, %10};\n" + : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3]) + : "r"(A[2]), "r"(A[3]), "r"(B[1]), "f"(C[0]), "f"(C[1]), "f"(C[2]), "f"(C[3])); + asm volatile( + "mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 " + "{%0, %1, %2, %3}," + "{%4, %5}," + "{%6}," + "{%7, %8, %9, %10};\n" + : "=f"(C[4]), "=f"(C[5]), "=f"(C[6]), "=f"(C[7]) + : "r"(A[0]), "r"(A[1]), "r"(B[2]), "f"(0.f), "f"(0.f), "f"(0.f), "f"(0.f)); + asm volatile( + "mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 " + "{%0, %1, %2, %3}," + "{%4, %5}," + "{%6}," + "{%7, %8, %9, %10};\n" + : "=f"(C[4]), "=f"(C[5]), "=f"(C[6]), "=f"(C[7]) + : "r"(A[2]), "r"(A[3]), "r"(B[3]), "f"(C[4]), "f"(C[5]), "f"(C[6]), "f"(C[7])); + } else { + asm volatile( + "mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 " + "{%0, %1, %2, %3}," + "{%4, %5}," + "{%6}," + "{%7, %8, %9, %10};\n" + : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3]) + : "r"(A[0]), "r"(A[1]), "r"(B[0]), "f"(C[0]), "f"(C[1]), "f"(C[2]), "f"(C[3])); + asm volatile( + "mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 " + "{%0, %1, %2, %3}," + "{%4, %5}," + "{%6}," + "{%7, %8, %9, %10};\n" + : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3]) + : "r"(A[2]), "r"(A[3]), "r"(B[1]), "f"(C[0]), "f"(C[1]), "f"(C[2]), "f"(C[3])); + asm volatile( + "mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 " + "{%0, %1, %2, %3}," + "{%4, %5}," + "{%6}," + "{%7, %8, %9, %10};\n" + : "=f"(C[4]), "=f"(C[5]), "=f"(C[6]), "=f"(C[7]) + : "r"(A[0]), "r"(A[1]), "r"(B[2]), "f"(C[4]), "f"(C[5]), "f"(C[6]), "f"(C[7])); + asm volatile( + "mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 " + "{%0, %1, %2, %3}," + "{%4, %5}," + "{%6}," + "{%7, %8, %9, %10};\n" + : "=f"(C[4]), "=f"(C[5]), "=f"(C[6]), "=f"(C[7]) + : "r"(A[2]), "r"(A[3]), "r"(B[3]), "f"(C[4]), "f"(C[5]), "f"(C[6]), "f"(C[7])); + } + } else { + FLASHINFER_RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction"); + } +#else + FLASHINFER_RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction"); +#endif +} + +/*! + * \brief Use mma instructions to compute rowsum. + */ +template +__device__ __forceinline__ void m16k32_rowsum_f8f8f32(float* d, DType* s) { + static_assert(sizeof(DType) == 1, "DType must be 8bit floating data type"); + uint32_t* s_u32 = (uint32_t*)(s); +#if defined(FLASHINFER_MMA_F8F8F32_M16N8K32_ENABLED) + if constexpr (std::is_same_v) { + asm volatile( + "{\n" + "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 " + "{%0, _, %1, _}," + "{%2, %3, %4, %5}," + "{%6, %7}," + "{%8, 0., %9, 0.};\n" + "}\n" + : "=f"(d[0]), "=f"(d[1]) + : "r"(s_u32[0]), "r"(s_u32[1]), "r"(s_u32[2]), "r"(s_u32[3]), "r"(943208504), + "r"(943208504), "f"(d[0]), "f"(d[1])); + } else { // e5m2 + asm volatile( + "{\n" + "mma.sync.aligned.m16n8k16.row.col.f32.e5m2.e5m2.f32 " + "{%0, _, %1, _}," + "{%2, %3, %4, %5}," + "{%6, %7}," + "{%8, 0., %9, 0.};\n" + "}\n" + : "=f"(d[0]), "=f"(d[1]) + : "r"(s_u32[0]), "r"(s_u32[1]), "r"(s_u32[2]), "r"(s_u32[3]), "r"(1010580540), + "r"(1010580540), "f"(d[0]), "f"(d[1])); + } +#else + FLASHINFER_RUNTIME_ASSERT( + "fp8 mma instruction is only available for sm89, PTX 8.4+ and CUDA 12.4+"); +#endif +} + +/*! + * \brief Use mma instructions to compute rowsum. + */ +template +__device__ __forceinline__ void m16k16_rowsum_f16f16f32(float* d, DType* s) { + static_assert(sizeof(DType) == 2, "DType must be 16bit floating data type"); + uint32_t* s_u32 = (uint32_t*)(s); +#if defined(FLASHINFER_MMA_F16F16F32_M16N8K16_ENABLED) + if constexpr (std::is_same_v) { + asm volatile( + "{\n" + "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 " + "{%0, _, %1, _}," + "{%2, %3, %4, %5}," + "{%6, %7}," + "{%8, 0., %9, 0.};\n" + "}\n" + : "=f"(d[0]), "=f"(d[1]) + : "r"(s_u32[0]), "r"(s_u32[1]), "r"(s_u32[2]), "r"(s_u32[3]), "r"(1006648320), + "r"(1006648320), "f"(d[0]), "f"(d[1])); + } else { + asm volatile( + "{\n" + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " + "{%0, _, %1, _}," + "{%2, %3, %4, %5}," + "{%6, %7}," + "{%8, 0., %9, 0.};\n" + "}\n" + : "=f"(d[0]), "=f"(d[1]) + : "r"(s_u32[0]), "r"(s_u32[1]), "r"(s_u32[2]), "r"(s_u32[3]), "r"(1065369472), + "r"(1065369472), "f"(d[0]), "f"(d[1])); + } +#elif defined(FLASHINFER_MMA_F16F16F32_M16N8K8_ENABLED) + if constexpr (std::is_same_v) { + asm volatile( + "{\n" + "mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 " + "{%0, _, %1, _}," + "{%2, %3}," + "{%4}," + "{%5, 0., %6, 0.};\n" + "}\n" + : "=f"(d[0]), "=f"(d[1]) + : "r"(s_u32[0]), "r"(s_u32[1]), "r"(1006648320), "f"(d[0]), "f"(d[1])); + asm volatile( + "{\n" + "mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 " + "{%0, _, %1, _}," + "{%2, %3}," + "{%4}," + "{%5, 0., %6, 0.};\n" + "}\n" + : "=f"(d[0]), "=f"(d[1]) + : "r"(s_u32[2]), "r"(s_u32[3]), "r"(1006648320), "f"(d[0]), "f"(d[1])); + } else { + FLASHINFER_RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction"); + } +#else + FLASHINFER_RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction"); +#endif +} + +/*! + * \brief Wrapper of two mma m16n8k16 instructions for row major and column major f16 matrix + * multiplication, accumulated in f16. + * \tparam mma_mode whether we are initializing the accumulator or updating it + * \param C pointer to the accumulator + * \param A pointer to the fragment of matrix A + * \param B pointer to the fragment of matrix B + */ +template +__device__ __forceinline__ void mma_sync_m16n16k16_row_col_f16f16f16(uint32_t* C, uint32_t* A, + uint32_t* B) { +#if defined(FLASHINFER_MMA_F16F16F16_M16N8K16_ENABLED) + if constexpr (mma_mode == MMAMode::kInit) { + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 " + "{%0, %1}," + "{%2, %3, %4, %5}," + "{%6, %7}," + "{%8, %9};\n" + : "=r"(C[0]), "=r"(C[1]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(0), "r"(0)); + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 " + "{%0, %1}," + "{%2, %3, %4, %5}," + "{%6, %7}," + "{%8, %9};\n" + : "=r"(C[2]), "=r"(C[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(0), "r"(0)); + } else { + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 " + "{%0, %1}," + "{%2, %3, %4, %5}," + "{%6, %7}," + "{%8, %9};\n" + : "=r"(C[0]), "=r"(C[1]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(C[0]), "r"(C[1])); + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 " + "{%0, %1}," + "{%2, %3, %4, %5}," + "{%6, %7}," + "{%8, %9};\n" + : "=r"(C[2]), "=r"(C[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(C[2]), "r"(C[3])); + } +#elif defined(FLASHINFER_MMA_F16F16F16_M16N8K8_ENABLED) + if constexpr (mma_mode == MMAMode::kInit) { + asm volatile( + "mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 " + "{%0, %1}," + "{%2, %3}," + "{%4}," + "{%5, %6};\n" + : "=r"(C[0]), "=r"(C[1]) + : "r"(A[0]), "r"(A[1]), "r"(B[0]), "r"(0), "r"(0)); + asm volatile( + "mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 " + "{%0, %1}," + "{%2, %3}," + "{%4}," + "{%5, %6};\n" + : "=r"(C[0]), "=r"(C[1]) + : "r"(A[2]), "r"(A[3]), "r"(B[1]), "r"(0), "r"(0)); + asm volatile( + "mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 " + "{%0, %1}," + "{%2, %3}," + "{%4}," + "{%5, %6};\n" + : "=r"(C[2]), "=r"(C[3]) + : "r"(A[0]), "r"(A[1]), "r"(B[2]), "r"(0), "r"(0)); + asm volatile( + "mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 " + "{%0, %1}," + "{%2, %3}," + "{%4}," + "{%5, %6};\n" + : "=r"(C[2]), "=r"(C[3]) + : "r"(A[2]), "r"(A[3]), "r"(B[3]), "r"(0), "r"(0)); + } else { + asm volatile( + "mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 " + "{%0, %1}," + "{%2, %3}," + "{%4}," + "{%5, %6};\n" + : "=r"(C[0]), "=r"(C[1]) + : "r"(A[0]), "r"(A[1]), "r"(B[0]), "r"(C[0]), "r"(C[1])); + asm volatile( + "mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 " + "{%0, %1}," + "{%2, %3}," + "{%4}," + "{%5, %6};\n" + : "=r"(C[0]), "=r"(C[1]) + : "r"(A[2]), "r"(A[3]), "r"(B[1]), "r"(C[0]), "r"(C[1])); + asm volatile( + "mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 " + "{%0, %1}," + "{%2, %3}," + "{%4}," + "{%5, %6};\n" + : "=r"(C[2]), "=r"(C[3]) + : "r"(A[0]), "r"(A[1]), "r"(B[2]), "r"(C[2]), "r"(C[3])); + asm volatile( + "mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 " + "{%0, %1}," + "{%2, %3}," + "{%4}," + "{%5, %6};\n" + : "=r"(C[2]), "=r"(C[3]) + : "r"(A[2]), "r"(A[3]), "r"(B[3]), "r"(C[2]), "r"(C[3])); + } +#else + FLASHINFER_RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction"); +#endif +} + +} // namespace mma + +} // namespace flashinfer + +#endif // FLASHINFER_MMA_CUH_ diff --git a/include/flashinfer/norm.cuh b/include/flashinfer/norm.cuh new file mode 100644 index 0000000000000000000000000000000000000000..f2c91138b36ca040aba0dc29e76b3b4517979bf2 --- /dev/null +++ b/include/flashinfer/norm.cuh @@ -0,0 +1,349 @@ +/* + * Copyright (c) 2024 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef FLASHINFER_NORM_CUH_ +#define FLASHINFER_NORM_CUH_ + +#include + +#include "flashinfer/utils.cuh" +#include "math.cuh" +#include "utils.cuh" +#include "vec_dtypes.cuh" + +namespace flashinfer { + +namespace norm { + +template +__global__ void RMSNormKernel(T* __restrict__ input, T* __restrict__ weight, T* __restrict__ output, + const uint32_t d, const uint32_t stride_input, + const uint32_t stride_output, float weight_bias, float eps) { + const uint32_t bx = blockIdx.x; + const uint32_t tx = threadIdx.x, ty = threadIdx.y; + constexpr uint32_t warp_size = 32; + const uint32_t num_warps = blockDim.y; + // NOTE(Zihao): it's guaranteed that num_warps should be smaller than 32 + const uint32_t thread_id = tx + ty * warp_size; + const uint32_t num_threads = num_warps * warp_size; + const uint32_t rounds = ceil_div(d, VEC_SIZE * num_threads); + extern __shared__ float smem[]; + + float sum_sq = 0.f; + +#if (__CUDACC_VER_MAJOR__ >= 12 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.wait;"); +#endif + + for (uint32_t i = 0; i < rounds; i++) { + vec_t input_vec; + input_vec.fill(0.f); + if ((i * num_threads + thread_id) * VEC_SIZE < d) { + input_vec.load(input + bx * stride_input + i * num_threads * VEC_SIZE + thread_id * VEC_SIZE); + } +#pragma unroll + for (uint32_t j = 0; j < VEC_SIZE; j++) { + sum_sq += float(input_vec[j]) * float(input_vec[j]); + } + } + + // first, warp reduce sum +#pragma unroll + for (uint32_t offset = warp_size / 2; offset > 0; offset /= 2) { + sum_sq += math::shfl_xor_sync(sum_sq, offset); + } + + smem[ty] = sum_sq; + __syncthreads(); + // then, cross warp reduce sum using only the first warp + if (ty == 0) { + sum_sq = (tx < num_warps) ? smem[tx] : 0.f; +#pragma unroll + for (uint32_t offset = warp_size / 2; offset > 0; offset /= 2) { + sum_sq += math::shfl_xor_sync(sum_sq, offset); + } + smem[0] = sum_sq; + } + __syncthreads(); + + float rms_rcp = math::rsqrt(smem[0] / float(d) + eps); + + for (uint32_t i = 0; i < rounds; i++) { + vec_t input_vec; + vec_t weight_vec; + vec_t output_vec; + input_vec.fill(0.f); + weight_vec.fill(0.f); + if ((i * num_threads + thread_id) * VEC_SIZE < d) { + input_vec.load(input + bx * stride_input + i * num_threads * VEC_SIZE + thread_id * VEC_SIZE); + weight_vec.load(weight + i * num_threads * VEC_SIZE + thread_id * VEC_SIZE); + } +#pragma unroll + for (uint32_t j = 0; j < VEC_SIZE; j++) { + output_vec[j] = float(input_vec[j]) * rms_rcp * (weight_bias + float(weight_vec[j])); + } + if ((i * num_threads + thread_id) * VEC_SIZE < d) { + output_vec.store(output + bx * stride_output + i * num_threads * VEC_SIZE + + thread_id * VEC_SIZE); + } + } +#if (__CUDACC_VER_MAJOR__ >= 12 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.launch_dependents;"); +#endif +} + +template +cudaError_t RMSNorm(T* input, T* weight, T* output, uint32_t batch_size, uint32_t d, + uint32_t stride_input, uint32_t stride_output, float eps = 1e-5, + bool enable_pdl = false, cudaStream_t stream = 0) { + const uint32_t vec_size = std::gcd(16 / sizeof(T), d); + + const uint32_t block_size = std::min(1024, d / vec_size); + const uint32_t num_warps = ceil_div(block_size, 32); + dim3 nblks(batch_size); + dim3 nthrs(32, num_warps); + const uint32_t smem_size = num_warps * sizeof(float); + float weight_bias = 0.f; + void* args[] = {&input, &weight, &output, &d, &stride_input, &stride_output, &weight_bias, &eps}; + + cudaLaunchConfig_t config; + config.gridDim = nblks; + config.blockDim = nthrs; + config.dynamicSmemBytes = smem_size; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = enable_pdl; + config.numAttrs = 1; + config.attrs = attrs; + + DISPATCH_ALIGNED_VEC_SIZE(vec_size, VEC_SIZE, { + auto kernel = RMSNormKernel; + FLASHINFER_CUDA_CALL( + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + FLASHINFER_CUDA_CALL(cudaLaunchKernelEx(&config, kernel, input, weight, output, d, stride_input, + stride_output, weight_bias, eps)); + }); + return cudaSuccess; +} + +template +__global__ void FusedAddRMSNormKernel(T* __restrict__ input, T* __restrict__ residual, + T* __restrict__ weight, const uint32_t d, + const uint32_t stride_input, const uint32_t stride_residual, + float weight_bias, float eps) { + const uint32_t bx = blockIdx.x; + const uint32_t tx = threadIdx.x, ty = threadIdx.y; + constexpr uint32_t warp_size = 32; + const uint32_t num_warps = blockDim.y; + const uint32_t thread_id = tx + ty * warp_size; + const uint32_t num_threads = num_warps * warp_size; + const uint32_t rounds = ceil_div(d, VEC_SIZE * num_threads); + extern __shared__ float smem[]; + float* smem_x = smem + ceil_div(num_warps, 4) * 4; + + float sum_sq = 0.f; +#if (__CUDACC_VER_MAJOR__ >= 12 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.wait;"); +#endif + + for (uint32_t i = 0; i < rounds; i++) { + vec_t input_vec; + input_vec.fill(0.f); + vec_t residual_vec; + residual_vec.fill(0.f); + vec_t x_vec; + x_vec.fill(0.f); + if ((i * num_threads + thread_id) * VEC_SIZE < d) { + input_vec.load(input + bx * stride_input + i * num_threads * VEC_SIZE + thread_id * VEC_SIZE); + residual_vec.load(residual + bx * stride_residual + i * num_threads * VEC_SIZE + + thread_id * VEC_SIZE); + } +#pragma unroll + for (uint32_t j = 0; j < VEC_SIZE; j++) { + float x = float(input_vec[j]); + x += float(residual_vec[j]); + sum_sq += x * x; + residual_vec[j] = (T)x; + x_vec[j] = x; + } + if ((i * num_threads + thread_id) * VEC_SIZE < d) { + residual_vec.store(residual + bx * stride_residual + i * num_threads * VEC_SIZE + + thread_id * VEC_SIZE); + x_vec.store(smem_x + i * num_threads * VEC_SIZE + thread_id * VEC_SIZE); + } + } + + // first, warp reduce sum +#pragma unroll + for (uint32_t offset = warp_size / 2; offset > 0; offset /= 2) { + sum_sq += math::shfl_xor_sync(sum_sq, offset); + } + + smem[ty] = sum_sq; + __syncthreads(); + // then, cross warp reduce sum using only the first warp + if (ty == 0) { + sum_sq = (tx < num_warps) ? smem[tx] : 0.f; +#pragma unroll + for (uint32_t offset = warp_size / 2; offset > 0; offset /= 2) { + sum_sq += math::shfl_xor_sync(sum_sq, offset); + } + smem[0] = sum_sq; + } + __syncthreads(); + + float rms_rcp = math::rsqrt(smem[0] / float(d) + eps); + + for (uint32_t i = 0; i < rounds; i++) { + vec_t input_vec; + vec_t weight_vec; + vec_t x_vec; + input_vec.fill(0.f); + weight_vec.fill(0.f); + x_vec.fill(0.f); + if ((i * num_threads + thread_id) * VEC_SIZE < d) { + weight_vec.load(weight + i * num_threads * VEC_SIZE + thread_id * VEC_SIZE); + x_vec.load(smem_x + i * num_threads * VEC_SIZE + thread_id * VEC_SIZE); + } +#pragma unroll + for (uint32_t j = 0; j < VEC_SIZE; j++) { + input_vec[j] = x_vec[j] * rms_rcp * (weight_bias + float(weight_vec[j])); + } + if ((i * num_threads + thread_id) * VEC_SIZE < d) { + input_vec.store(input + bx * stride_input + i * num_threads * VEC_SIZE + + thread_id * VEC_SIZE); + } + } +#if (__CUDACC_VER_MAJOR__ >= 12 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.launch_dependents;"); +#endif +} + +template +cudaError_t FusedAddRMSNorm(T* input, T* residual, T* weight, uint32_t batch_size, uint32_t d, + uint32_t stride_input, uint32_t stride_residual, float eps = 1e-5, + bool enable_pdl = false, cudaStream_t stream = 0) { + const uint32_t vec_size = std::gcd(16 / sizeof(T), d); + + const uint32_t block_size = std::min(1024, d / vec_size); + const uint32_t num_warps = ceil_div(block_size, 32); + dim3 nblks(batch_size); + dim3 nthrs(32, num_warps); + const uint32_t smem_size = (ceil_div(num_warps, 4) * 4 + d) * sizeof(float); + float weight_bias = 0.f; + void* args[] = {&input, &residual, &weight, &d, + &stride_input, &stride_residual, &weight_bias, &eps}; + + cudaLaunchConfig_t config; + config.gridDim = nblks; + config.blockDim = nthrs; + config.dynamicSmemBytes = smem_size; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = enable_pdl; + config.numAttrs = 1; + config.attrs = attrs; + + DISPATCH_ALIGNED_VEC_SIZE(vec_size, VEC_SIZE, { + auto kernel = FusedAddRMSNormKernel; + FLASHINFER_CUDA_CALL( + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + FLASHINFER_CUDA_CALL(cudaLaunchKernelEx(&config, kernel, input, residual, weight, d, + stride_input, stride_residual, weight_bias, eps)); + }); + + return cudaSuccess; +} + +template +cudaError_t GemmaRMSNorm(T* input, T* weight, T* output, uint32_t batch_size, uint32_t d, + uint32_t stride_input, uint32_t stride_output, float eps = 1e-5, + bool enable_pdl = false, cudaStream_t stream = 0) { + const uint32_t vec_size = std::gcd(16 / sizeof(T), d); + + const uint32_t block_size = std::min(1024, d / vec_size); + const uint32_t num_warps = ceil_div(block_size, 32); + dim3 nblks(batch_size); + dim3 nthrs(32, num_warps); + const uint32_t smem_size = num_warps * sizeof(float); + float weight_bias = 1.f; + void* args[] = {&input, &weight, &output, &d, &stride_input, &stride_output, &weight_bias, &eps}; + + cudaLaunchConfig_t config; + config.gridDim = nblks; + config.blockDim = nthrs; + config.dynamicSmemBytes = smem_size; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = enable_pdl; + config.numAttrs = 1; + config.attrs = attrs; + + DISPATCH_ALIGNED_VEC_SIZE(vec_size, VEC_SIZE, { + auto kernel = RMSNormKernel; + FLASHINFER_CUDA_CALL( + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + FLASHINFER_CUDA_CALL(cudaLaunchKernelEx(&config, kernel, input, weight, output, d, stride_input, + stride_output, weight_bias, eps)); + }); + return cudaSuccess; +} + +template +cudaError_t GemmaFusedAddRMSNorm(T* input, T* residual, T* weight, uint32_t batch_size, uint32_t d, + uint32_t stride_input, uint32_t stride_residual, float eps = 1e-5, + bool enable_pdl = false, cudaStream_t stream = 0) { + const uint32_t vec_size = std::gcd(16 / sizeof(T), d); + + const uint32_t block_size = std::min(1024, d / vec_size); + const uint32_t num_warps = ceil_div(block_size, 32); + dim3 nblks(batch_size); + dim3 nthrs(32, num_warps); + // NOTE(Zihao): use ceil_div(num_warps, 4) * 4 for address alignment to 16 bytes + const uint32_t smem_size = (ceil_div(num_warps, 4) * 4 + d) * sizeof(float); + float weight_bias = 1.f; + void* args[] = {&input, &residual, &weight, &d, + &stride_input, &stride_residual, &weight_bias, &eps}; + + cudaLaunchConfig_t config; + config.gridDim = nblks; + config.blockDim = nthrs; + config.dynamicSmemBytes = smem_size; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = enable_pdl; + config.numAttrs = 1; + config.attrs = attrs; + + DISPATCH_ALIGNED_VEC_SIZE(vec_size, VEC_SIZE, { + auto kernel = FusedAddRMSNormKernel; + FLASHINFER_CUDA_CALL( + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + FLASHINFER_CUDA_CALL(cudaLaunchKernelEx(&config, kernel, input, residual, weight, d, + stride_input, stride_residual, weight_bias, eps)); + }); + + return cudaSuccess; +} + +} // namespace norm + +} // namespace flashinfer + +#endif // FLASHINFER_NORM_CUH_ diff --git a/include/flashinfer/page.cuh b/include/flashinfer/page.cuh new file mode 100644 index 0000000000000000000000000000000000000000..1f5d328da87284fe3aee5c6f63043510d070fb0f --- /dev/null +++ b/include/flashinfer/page.cuh @@ -0,0 +1,646 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef FLASHINFER_PAGE_CUH_ +#define FLASHINFER_PAGE_CUH_ + +#include + +#include + +#include "exception.h" +#include "fastdiv.cuh" +#include "layout.cuh" +#include "utils.cuh" +#include "vec_dtypes.cuh" + +namespace flashinfer { + +/*! + * \brief Paged key-value cache + * \tparam layout The layout of last 3 dimensions in KV-Cache. + * \tparam DType The data type of the key-value cache + * \tparam IdType The index data type of the kv-cache + */ +template +struct paged_kv_t { + uint_fastdiv page_size; + uint32_t num_heads; + uint32_t head_dim; + uint32_t batch_size; + uint32_t stride_page; + uint32_t stride_n; + uint32_t stride_h; + + // Internal layout: + // [max_num_pages, num_heads, page_size, head_dim] if layout == HND + // [max_num_pages, page_size, num_heads, head_dim] if layout == NHD + DType* k_data; + DType* v_data; + IdType* indices; + + // [batch_size + 1] The page indptr array, with the first element 0, the last element nnz_pages + IdType* indptr; + // [batch_size] The offset of the last page for each request in the batch + IdType* last_page_len; + // [batch_size] The start position of each request in the batch. + IdType* rope_pos_offset; + + /*! + * \brief Construct an empty paged key-value cache + */ + __host__ __device__ __forceinline__ paged_kv_t() + : num_heads(0), + page_size(), + head_dim(0), + batch_size(0), + stride_page(0), + stride_n(0), + stride_h(0), + k_data(nullptr), + v_data(nullptr), + indices(nullptr), + indptr(nullptr), + last_page_len(nullptr), + rope_pos_offset(nullptr) {} + + /*! + * \brief Construct a paged key-value cache + * \param num_heads The number of heads + * \param page_size The size of each page + * \param head_dim The dimension of each head + * \param batch_size The batch size + * \param layout The layout of last 3 dimensions in KV-Cache. + * \param k_data The start pointer of key cache, k_cache should be contiguous + * \param v_data The start pointer of value cache, v_cache should be contiguous + * \param indices The page indices array + * \param indptr The page indptr array + * \param last_page_len The offset of the last page for each request in the batch + * \param rope_pos_offset The start position of each request in the batch. + */ + __host__ __forceinline__ paged_kv_t(uint32_t num_heads, uint32_t page_size, uint32_t head_dim, + uint32_t batch_size, QKVLayout layout, DType* k_data, + DType* v_data, IdType* indices, IdType* indptr, + IdType* last_page_len, IdType* rope_pos_offset = nullptr) + : num_heads(num_heads), + page_size(page_size), + head_dim(head_dim), + batch_size(batch_size), + indices(indices), + indptr(indptr), + last_page_len(last_page_len), + rope_pos_offset(rope_pos_offset) { + stride_page = num_heads * page_size * head_dim; + this->k_data = k_data; + this->v_data = v_data; + stride_n = layout == QKVLayout::kHND ? head_dim : num_heads * head_dim; + stride_h = layout == QKVLayout::kHND ? page_size * head_dim : head_dim; + } + + /*! + * \brief Construct a paged key-value cache with custom kv-cache strides + * \param num_heads The number of heads + * \param page_size The size of each page + * \param head_dim The dimension of each head + * \param batch_size The batch size + * \param layout The layout of last 3 dimensions in KV-Cache. + * \param k_data The start pointer of key cache, k_cache doesn't have to be contiguous + * \param v_data The start pointer of value cache, v_cache doesn't have to be contiguous + * \param kv_strides custom strides of each dimensions of k_data and v_data + * \param indices The page indices array + * \param indptr The page indptr array + * \param last_page_len The offset of the last page for each request in the batch + * \param rope_pos_offset The start position of each request in the batch. + */ + __host__ __forceinline__ paged_kv_t(uint32_t num_heads, uint32_t page_size, uint32_t head_dim, + uint32_t batch_size, QKVLayout layout, DType* k_data, + DType* v_data, const int64_t* kv_strides, IdType* indices, + IdType* indptr, IdType* last_page_len, + IdType* rope_pos_offset = nullptr) + : num_heads(num_heads), + page_size(page_size), + head_dim(head_dim), + batch_size(batch_size), + indices(indices), + indptr(indptr), + last_page_len(last_page_len), + rope_pos_offset(rope_pos_offset) { + stride_page = kv_strides[0]; + this->k_data = k_data; + this->v_data = v_data; + stride_n = layout == QKVLayout::kHND ? kv_strides[2] : kv_strides[1]; + stride_h = layout == QKVLayout::kHND ? kv_strides[1] : kv_strides[2]; + } + + __host__ __device__ __forceinline__ uint32_t get_length(uint32_t batch_idx) const { + if (indptr[batch_idx + 1] == indptr[batch_idx]) { + return 0; + } + return (indptr[batch_idx + 1] - indptr[batch_idx] - 1) * page_size + last_page_len[batch_idx]; + } + + /*! + * \brief Compute the offset of element in the allocated buffer. + * \param page_idx The page index + * \param head_idx The head index + * \param entry_idx The page entry index + * \param feat_idx The feature index + */ + __host__ __device__ __forceinline__ size_t get_elem_offset(size_t page_idx, size_t head_idx, + size_t entry_idx, + size_t feat_idx) const { + return page_idx * stride_page + head_idx * stride_h + entry_idx * stride_n + feat_idx; + } + + /*! + * \brief Compute the offset of element inside the page. + * \param head_idx The head index + * \param entry_idx The page entry index + * \param feat_idx The feature index + */ + __host__ __device__ __forceinline__ size_t get_elem_offset_in_page(size_t head_idx, + size_t entry_idx, + size_t feat_idx) const { + return head_idx * stride_h + entry_idx * stride_n + feat_idx; + } + + __device__ __forceinline__ DType* get_k_ptr(IdType page_iter, uint32_t head_idx, + uint32_t entry_idx, uint32_t feat_idx) const { + return k_data + get_elem_offset(__ldg(indices + page_iter), head_idx, entry_idx, feat_idx); + } + + __device__ __forceinline__ size_t protective_get_kv_offset(IdType page_iter, uint32_t head_idx, + uint32_t entry_idx, uint32_t feat_idx, + IdType last_indptr) const { + if (page_iter < last_indptr) { + return get_elem_offset(__ldg(indices + page_iter), head_idx, entry_idx, feat_idx); + } else { + return 0; + } + } + + __device__ __forceinline__ DType* protective_get_k_ptr(IdType page_iter, uint32_t head_idx, + uint32_t entry_idx, uint32_t feat_idx, + IdType last_indptr) const { + return k_data + protective_get_kv_offset(page_iter, head_idx, entry_idx, feat_idx, last_indptr); + } + + __device__ __forceinline__ DType* get_v_ptr(IdType page_iter, uint32_t head_idx, + uint32_t entry_idx, uint32_t feat_idx) const { + return v_data + get_elem_offset(__ldg(indices + page_iter), head_idx, entry_idx, feat_idx); + } + + __device__ __forceinline__ DType* protective_get_v_ptr(IdType page_iter, uint32_t head_idx, + uint32_t entry_idx, uint32_t feat_idx, + IdType last_indptr) const { + return v_data + protective_get_kv_offset(page_iter, head_idx, entry_idx, feat_idx, last_indptr); + } +}; + +/*! + * \brief CUDA kernel to append new keys/values to the paged key-value cache in the decode phase + * \tparam head_dim The dimension of each head + * \tparam vec_size The vector size used in the kernel + * \tparam DType The data type of the key-value cache + * \tparam IdType The index data type of the kv-cache + * \param paged_kv The paged key-value cache + * \param key The key to be appended + * \param value The value to be appended + */ +template +__global__ void AppendPagedKVCacheDecodeKernel(paged_kv_t paged_kv, + DType* __restrict__ key, DType* __restrict__ value) { + uint32_t tx = threadIdx.x, ty = threadIdx.y; + uint32_t num_heads = paged_kv.num_heads; + uint32_t batch_idx = blockIdx.x; + uint32_t head_idx = ty; + + uint32_t seq_len = + (paged_kv.indptr[batch_idx + 1] - paged_kv.indptr[batch_idx] - 1) * paged_kv.page_size + + paged_kv.last_page_len[batch_idx]; + + uint32_t page_iter = paged_kv.indptr[batch_idx] + (seq_len - 1) / paged_kv.page_size; + uint32_t entry_idx = (seq_len - 1) % paged_kv.page_size; + + DType* k_ptr = paged_kv.get_k_ptr(page_iter, head_idx, entry_idx, tx * vec_size); + DType* v_ptr = paged_kv.get_v_ptr(page_iter, head_idx, entry_idx, tx * vec_size); + vec_t::memcpy( + k_ptr, key + (batch_idx * num_heads + head_idx) * head_dim + tx * vec_size); + + vec_t::memcpy( + v_ptr, value + (batch_idx * num_heads + head_idx) * head_dim + tx * vec_size); +} + +/*! + * \brief CUDA kernel to append new keys/values to the paged key-value cache in the prefill phase + * \tparam head_dim The dimension of each head + * \tparam vec_size The vector size used in the kernel + * \tparam DType The data type of the key-value cache + * \tparam IdType The index data type of the kv-cache + * \param paged_kv The paged key-value cache + * \param key The key to be appended + * \param value The value to be appended + * \param batch_indices The batch indices of elements to be appended + * \param positions The positions of elements to be appended + */ +template +__global__ void AppendPagedKVCacheKernel(paged_kv_t paged_kv, + DType* __restrict__ append_key, + DType* __restrict__ append_value, + IdType* __restrict__ batch_indices, + IdType* __restrict__ positions, uint32_t nnz, + size_t append_k_stride_n, size_t append_k_stride_h, + size_t append_v_stride_n, size_t append_v_stride_h) { + uint32_t tx = threadIdx.x, ty = threadIdx.y; + uint32_t num_heads = paged_kv.num_heads; + uint32_t head_idx = ty; + uint32_t cta_id = blockIdx.x; + uint32_t num_ctas = gridDim.x; + +#pragma unroll 4 + for (uint32_t i = cta_id; i < nnz; i += num_ctas) { + uint32_t page_iter, entry_idx; + paged_kv.page_size.divmod(paged_kv.indptr[batch_indices[i]] * paged_kv.page_size + positions[i], + page_iter, entry_idx); + DType* k_ptr = paged_kv.get_k_ptr(page_iter, head_idx, entry_idx, tx * vec_size); + DType* v_ptr = paged_kv.get_v_ptr(page_iter, head_idx, entry_idx, tx * vec_size); + vec_t::memcpy( + k_ptr, append_key + i * append_k_stride_n + head_idx * append_k_stride_h + tx * vec_size); + vec_t::memcpy( + v_ptr, append_value + i * append_v_stride_n + head_idx * append_v_stride_h + tx * vec_size); + } +} + +template +__global__ void BlockSparseIndicesToVectorSparseOffsetsKernel( + IdType* __restrict__ block_sparse_indices, IdType* __restrict__ block_sparse_indptr, + IdType* __restrict__ vector_sparse_offsets, IdType* __restrict__ vector_sparse_indptr, + IdType* __restrict__ kv_lens, const uint32_t stride_block, const uint32_t stride_n, + const uint32_t batch_size, const uint_fastdiv block_size) { +#pragma unroll 1 + for (int b = blockIdx.x; b < batch_size; ++b) { +#pragma unroll 2 + for (int pos = threadIdx.x; pos < kv_lens[b]; pos += blockDim.x) { + uint32_t q, r; + block_size.divmod(pos, q, r); + vector_sparse_offsets[vector_sparse_indptr[b] + pos] = + block_sparse_indices[block_sparse_indptr[b] + q] * stride_block + r * stride_n; + } + } +} + +template +cudaError_t BlockSparseIndicesToVectorSparseOffset( + IdType* block_sparse_indices, IdType* block_sparse_indptr, IdType* vector_sparse_offsets, + IdType* vector_sparse_indptr, IdType* kv_lens, const int64_t stride_block, + const int64_t stride_n, const int64_t batch_size, const uint32_t block_size, + cudaStream_t stream = nullptr) { + int dev_id = 0; + int num_sms = 0; + FLASHINFER_CUDA_CALL(cudaGetDevice(&dev_id)); + FLASHINFER_CUDA_CALL(cudaDeviceGetAttribute(&num_sms, cudaDevAttrMultiProcessorCount, dev_id)); + + uint32_t num_threads = 512; + + uint_fastdiv block_size_fastdiv(block_size); + + auto kernel = BlockSparseIndicesToVectorSparseOffsetsKernel; + void* args[] = {(void*)&block_sparse_indices, + (void*)&block_sparse_indptr, + (void*)&vector_sparse_offsets, + (void*)&vector_sparse_indptr, + (void*)&kv_lens, + (void*)&stride_block, + (void*)&stride_n, + (void*)&batch_size, + (void*)&block_size_fastdiv}; + + FLASHINFER_CUDA_CALL(cudaLaunchKernel((void*)kernel, num_sms, num_threads, args, 0, stream)); + + return cudaSuccess; +} + +/*! + * \brief Append new keys/values to the paged key-value cache in the decode phase + * \tparam DType The data type of the key-value cache + * \tparam IdType The index data type of the kv-cache + * \param paged_kv The paged key-value cache + * \param key The key to be appended + * \param value The value to be appended + * \param stream The CUDA stream to execute kernels. + * \return status Indicates whether CUDA calls are successful + */ +template +cudaError_t AppendPagedKVCacheDecode(paged_kv_t paged_kv, DType* key, DType* value, + cudaStream_t stream = nullptr) { + uint32_t head_dim = paged_kv.head_dim; + uint32_t batch_size = paged_kv.batch_size; + uint32_t num_heads = paged_kv.num_heads; + DISPATCH_HEAD_DIM(head_dim, HEAD_DIM, { + constexpr uint32_t vec_size = std::max(16 / sizeof(DType), HEAD_DIM / 32); + uint32_t bdx = HEAD_DIM / vec_size; + uint32_t bdy = num_heads; + // NOTE(Zihao): could be slow for small batch size, will optimize later + dim3 nblks(batch_size); + dim3 nthrs(bdx, bdy); + auto kernel = AppendPagedKVCacheDecodeKernel; + void* args[] = {(void*)&paged_kv, (void*)&key, (void*)&value}; + FLASHINFER_CUDA_CALL(cudaLaunchKernel((void*)kernel, nblks, nthrs, args, 0, stream)); + }); + return cudaSuccess; +} + +/*! + * \brief Append new keys/values to the paged key-value cache + * \tparam layout The layout of last 3 dimension in KV-Cache + * \tparam DType The data type of the key-value cache + * \tparam IdType The index data type of the kv-cache + * \param paged_kv The paged key-value cache + * \param key The key to be appended + * \param value The value to be appended + * \param append_indptr The indptr array of the appended ragged tensor + * \param stream The CUDA stream to execute kernels. + * \return status Indicates whether CUDA calls are successful + */ +template +cudaError_t AppendPagedKVCache(paged_kv_t paged_kv, DType* append_key, + DType* append_value, IdType* batch_indices, IdType* positions, + uint32_t nnz, size_t append_k_stride_n, size_t append_k_stride_h, + size_t append_v_stride_n, size_t append_v_stride_h, + cudaStream_t stream = nullptr) { + uint32_t head_dim = paged_kv.head_dim; + uint32_t num_heads = paged_kv.num_heads; + int dev_id = 0; + int num_sms = 0; + int num_blocks_per_sm = 0; + FLASHINFER_CUDA_CALL(cudaGetDevice(&dev_id)); + FLASHINFER_CUDA_CALL(cudaDeviceGetAttribute(&num_sms, cudaDevAttrMultiProcessorCount, dev_id)); + + DISPATCH_HEAD_DIM(head_dim, HEAD_DIM, { + constexpr uint32_t vec_size = std::max(16 / sizeof(DType), HEAD_DIM / 32); + uint32_t bdx = HEAD_DIM / vec_size; + uint32_t bdy = num_heads; + uint32_t num_threads = bdx * bdy; + uint32_t smem_size = 0; + auto kernel = AppendPagedKVCacheKernel; + FLASHINFER_CUDA_CALL(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&num_blocks_per_sm, kernel, + num_threads, smem_size)); + num_blocks_per_sm = min(num_blocks_per_sm, ceil_div(int(nnz), num_sms)); + dim3 nblks(num_blocks_per_sm * num_sms); + dim3 nthrs(bdx, bdy); + + void* args[] = {(void*)&paged_kv, (void*)&append_key, (void*)&append_value, + (void*)&batch_indices, (void*)&positions, (void*)&nnz, + (void*)&append_k_stride_n, (void*)&append_k_stride_h, (void*)&append_v_stride_n, + (void*)&append_v_stride_h}; + FLASHINFER_CUDA_CALL(cudaLaunchKernel((void*)kernel, nblks, nthrs, args, 0, stream)); + }); + return cudaSuccess; +} + +template +struct paged_kv_mla_t { + uint_fastdiv page_size; + uint32_t head_dim_ckv; + uint32_t head_dim_kpe; + uint32_t batch_size; + uint32_t stride_page_ckv; + uint32_t stride_page_kpe; + uint32_t stride_n_ckv; + uint32_t stride_n_kpe; + + // Internal layout: + // [max_num_pages, page_size, head_dim] + DType* ckv_data; + DType* kpe_data; + IdType* indices; + + // [batch_size + 1] The page indptr array, with the first element 0, the last element nnz_pages + IdType* indptr; + // [batch_size] The offset of the last page for each request in the batch + IdType* last_page_len; + // [batch_size] The start position of each request in the batch. + IdType* rope_pos_offset; + + /*! + * \brief Construct an empty paged key-value cache + */ + __host__ __device__ __forceinline__ paged_kv_mla_t() + : head_dim_ckv(0), + head_dim_kpe(0), + batch_size(0), + stride_page_ckv(0), + stride_page_kpe(0), + stride_n_ckv(0), + stride_n_kpe(0), + ckv_data(nullptr), + kpe_data(nullptr), + indices(nullptr), + indptr(nullptr), + last_page_len(nullptr), + rope_pos_offset(nullptr) {} + + /*! + * \brief Construct a paged mla kv cache + * \param page_size The size of each page + * \param head_dim_compressed_kv The dimension of compressed-kv + * \param head_dim_kpe The dimension of k-pe + * \param batch_size The batch size + * \param compressed_kv_data The start pointer of compressed-kv cache, cache should be contiguous + * \param kpe_data The start pointer of k-pe cache, cache should be contiguous + * \param indices The page indices array + * \param indptr The page indptr array + * \param last_page_len The offset of the last page for each request in the batch + * \param rope_pos_offset The start position of each request in the batch. + */ + __host__ __forceinline__ paged_kv_mla_t(uint32_t page_size, uint32_t head_dim_compressed_kv, + uint32_t head_dim_kpe, uint32_t batch_size, + DType* compressed_kv_data, DType* kpe_data, + IdType* indices, IdType* indptr, IdType* last_page_len, + IdType* rope_pos_offset = nullptr) + : page_size(page_size), + head_dim_ckv(head_dim_compressed_kv), + head_dim_kpe(head_dim_kpe), + batch_size(batch_size), + ckv_data(compressed_kv_data), + kpe_data(kpe_data), + indices(indices), + indptr(indptr), + last_page_len(last_page_len), + rope_pos_offset(rope_pos_offset) { + stride_page_ckv = page_size * head_dim_ckv; + stride_n_ckv = head_dim_ckv; + stride_page_kpe = page_size * head_dim_kpe; + stride_n_kpe = head_dim_kpe; + } + + /*! + * \brief Construct a paged key-value cache with custom kv-cache strides + * \param page_size The size of each page + * \param head_dim_compressed_kv The dimension of compressed-kv + * \param head_dim_kpe The dimension of k-pe + * \param batch_size The batch size + * \param compressed_kv_data The start pointer of compressed-kv cache, cache should be contiguous + * \param compressed_kv_strides custom strides of each dimensions of compressed-kv cache + * \param kpe_data The start pointer of k-pe cache, cache should be contiguous + * \param kpe_strides custom strides of each dimensions of k-pe cache + * \param indices The page indices array + * \param indptr The page indptr array + * \param last_page_len The offset of the last page for each request in the batch + * \param rope_pos_offset The start position of each request in the batch. + */ + __host__ __forceinline__ paged_kv_mla_t(uint32_t page_size, uint32_t head_dim_compressed_kv, + uint32_t head_dim_kpe, uint32_t batch_size, + DType* compressed_kv_data, + const int64_t* compressed_kv_strides, DType* kpe_data, + const int64_t* kpe_strides, IdType* indices, + IdType* indptr, IdType* last_page_len, + IdType* rope_pos_offset = nullptr) + : page_size(page_size), + head_dim_ckv(head_dim_compressed_kv), + head_dim_kpe(head_dim_kpe), + batch_size(batch_size), + ckv_data(compressed_kv_data), + kpe_data(kpe_data), + indices(indices), + indptr(indptr), + last_page_len(last_page_len), + rope_pos_offset(rope_pos_offset) { + stride_page_ckv = compressed_kv_strides[0]; + stride_n_ckv = compressed_kv_strides[1]; + stride_page_kpe = kpe_strides[0]; + stride_n_kpe = kpe_strides[1]; + } + + __host__ __device__ __forceinline__ uint32_t get_length(uint32_t batch_idx) const { + if (indptr[batch_idx + 1] == indptr[batch_idx]) { + return 0; + } + return (indptr[batch_idx + 1] - indptr[batch_idx] - 1) * page_size + last_page_len[batch_idx]; + } + + __host__ __device__ __forceinline__ size_t get_elem_offset_ckv(size_t page_idx, size_t entry_idx, + size_t feat_idx) const { + return page_idx * stride_page_ckv + entry_idx * stride_n_ckv + feat_idx; + } + + __device__ __forceinline__ size_t protective_get_offset_ckv(IdType page_iter, uint32_t entry_idx, + uint32_t feat_idx, + IdType last_indptr) const { + if (page_iter < last_indptr) { + return get_elem_offset_ckv(__ldg(indices + page_iter), entry_idx, feat_idx); + } else { + return 0; + } + } + + __host__ __device__ __forceinline__ size_t get_elem_offset_kpe(size_t page_idx, size_t entry_idx, + size_t feat_idx) const { + return page_idx * stride_page_kpe + entry_idx * stride_n_kpe + feat_idx; + } + + __device__ __forceinline__ size_t protective_get_offset_kpe(IdType page_iter, uint32_t entry_idx, + uint32_t feat_idx, + IdType last_indptr) const { + if (page_iter < last_indptr) { + return get_elem_offset_kpe(__ldg(indices + page_iter), entry_idx, feat_idx); + } else { + return 0; + } + } + + __device__ __forceinline__ DType* get_ckv_ptr(size_t page_idx, size_t entry_idx, + size_t feat_idx) const { + return ckv_data + get_elem_offset_ckv(__ldg(indices + page_idx), entry_idx, feat_idx); + } + + __device__ __forceinline__ DType* get_kpe_ptr(size_t page_idx, size_t entry_idx, + size_t feat_idx) const { + return kpe_data + get_elem_offset_kpe(__ldg(indices + page_idx), entry_idx, feat_idx); + } +}; + +template +__global__ void AppendPagedKVMlaCacheKernel(paged_kv_mla_t paged_kv_mla, + DType* __restrict__ append_ckv, + DType* __restrict__ append_kpe, + IdType* __restrict__ batch_indices, + IdType* __restrict__ positions, uint32_t nnz, + size_t append_ckv_stride_n, + size_t append_kpe_stride_n) { + uint32_t tx = threadIdx.x; + uint32_t cta_id = blockIdx.x; + uint32_t num_ctas = gridDim.x; + +#pragma unroll 4 + for (uint32_t i = cta_id; i < nnz; i += num_ctas) { + uint32_t page_iter, entry_idx; + paged_kv_mla.page_size.divmod( + paged_kv_mla.indptr[batch_indices[i]] * paged_kv_mla.page_size + positions[i], page_iter, + entry_idx); + DType* ckv_ptr = paged_kv_mla.get_ckv_ptr(page_iter, entry_idx, tx * vec_size); + vec_t::memcpy(ckv_ptr, append_ckv + i * append_ckv_stride_n + tx * vec_size); + + if (tx * vec_size < head_dim_kpe) { + DType* kpe_ptr = paged_kv_mla.get_kpe_ptr(page_iter, entry_idx, tx * vec_size); + vec_t::memcpy(kpe_ptr, append_kpe + i * append_kpe_stride_n + tx * vec_size); + } + } +} + +template +cudaError_t AppendPagedKVMlaCache(paged_kv_mla_t paged_kv, DType* append_ckv, + DType* append_kpe, IdType* batch_indices, IdType* positions, + uint32_t nnz, size_t append_ckv_stride_n, + size_t append_kpe_stride_n, cudaStream_t stream = nullptr) { + int dev_id = 0; + int num_sms = 0; + int num_blocks_per_sm = 0; + FLASHINFER_CUDA_CALL(cudaGetDevice(&dev_id)); + FLASHINFER_CUDA_CALL(cudaDeviceGetAttribute(&num_sms, cudaDevAttrMultiProcessorCount, dev_id)); + + uint32_t head_dim_ckv = paged_kv.head_dim_ckv; + uint32_t head_dim_kpe = paged_kv.head_dim_kpe; + constexpr uint32_t HEAD_CKV_DIM = 512; + constexpr uint32_t HEAD_KPE_DIM = 64; + FLASHINFER_CHECK(head_dim_ckv == HEAD_CKV_DIM, "head_dim_ckv must be equal to 512"); + FLASHINFER_CHECK(head_dim_kpe == HEAD_KPE_DIM, "head_dim_kpe must be equal to 64"); + constexpr uint32_t vec_size = 2; + + uint32_t bdx = HEAD_CKV_DIM / vec_size; + uint32_t num_threads = bdx; + uint32_t smem_size = 0; + auto kernel = AppendPagedKVMlaCacheKernel; + FLASHINFER_CUDA_CALL(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&num_blocks_per_sm, kernel, + num_threads, smem_size)); + num_blocks_per_sm = min(num_blocks_per_sm, ceil_div(int(nnz), num_sms)); + dim3 nblks(num_blocks_per_sm * num_sms); + dim3 nthrs(bdx); + void* args[] = {(void*)&paged_kv, + (void*)&append_ckv, + (void*)&append_kpe, + (void*)&batch_indices, + (void*)&positions, + (void*)&nnz, + (void*)&append_ckv_stride_n, + (void*)&append_kpe_stride_n}; + FLASHINFER_CUDA_CALL(cudaLaunchKernel((void*)kernel, nblks, nthrs, args, 0, stream)); + return cudaSuccess; +} + +} // namespace flashinfer + +#endif // FLAHSINFER_PAGE_CUH_ diff --git a/include/flashinfer/permuted_smem.cuh b/include/flashinfer/permuted_smem.cuh new file mode 100644 index 0000000000000000000000000000000000000000..a63283ebdc83c5dda3fe83d5ba187f0508c146e8 --- /dev/null +++ b/include/flashinfer/permuted_smem.cuh @@ -0,0 +1,184 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef FLASHINFER_PERMUTED_SMEM_CUH_ +#define FLASHINFER_PERMUTED_SMEM_CUH_ + +#include +#include +#include + +#include + +#include "cp_async.cuh" +#include "mma.cuh" + +namespace flashinfer { + +enum class SwizzleMode { + k64B, + k128B, +}; + +// Use 128bit as the granularity to fetch/store data per thread to maximize memory bandwidth +using b128_t = uint4; + +/*! + * \brief Compute the number of elements that can be stored in a b128_t. + * \tparam T The data type of the elements. + */ +template +constexpr __host__ __device__ __forceinline__ uint32_t upcast_size() { + return sizeof(b128_t) / sizeof(T); +} + +template +__device__ __forceinline__ uint32_t get_permuted_offset(uint32_t i, uint32_t j) { + if constexpr (swizzle_mode == SwizzleMode::k128B) { + return i * stride + (j ^ (i % 8)); + } else { + // swizzle_mode == SwizzleMode::k64B + return i * stride + (j ^ ((i / 2) % 4)); + } +} + +/*! + * \brief The shared memory wrapper. + */ +template +struct smem_t { + // The base pointer. + b128_t* base; + __device__ __forceinline__ smem_t() : base(nullptr) {} + template + __device__ __forceinline__ smem_t(T* base) : base((b128_t*)base) {} + + /*! + * \brief Compute the element offset given coordinates in a permuted shared memory. + * \tparam stride The stride (in terms of b128_t's) in the permuted shared memory. + * \param i The row index. + * \param j The column index. + */ + template + static __device__ __forceinline__ uint32_t get_permuted_offset(uint32_t i, uint32_t j) { + if constexpr (swizzle_mode == SwizzleMode::k128B) { + return i * stride + (j ^ (i % 8)); + } else { + // swizzle_mode == SwizzleMode::k64B + static_assert(stride == 4); + return i * stride + (j ^ ((i / 2) % 4)); + } + } + + template + static __device__ __forceinline__ uint32_t advance_offset_by_column(uint32_t offset, + uint32_t step_idx) { + if constexpr (swizzle_mode == SwizzleMode::k128B) { + static_assert(step_size == 2 || step_size == 4 || step_size % 8 == 0, + "Unsupported step size"); + if constexpr (step_size == 2) { + return (offset ^ (0x2 + (0x4 * (step_idx % 2 == 1)))) + (step_idx % 4 == 3) * 8; + } else if constexpr (step_size == 4) { + return (offset ^ 0x4) + (step_idx % 2 == 1) * 8; + } else { + // step_size % 8 == 0 + return offset + step_size; + } + } else { + // swizzle_mode == SwizzleMode::k64B + static_assert(step_size == 2, "Unsupported step size"); + return (offset ^ 0x2) + (step_idx % 2 == 1) * 4; + } + } + + template + static __device__ __forceinline__ uint32_t advance_offset_by_row(uint32_t offset) { + if constexpr (swizzle_mode == SwizzleMode::k128B) { + static_assert(step_size == 4 || step_size % 8 == 0, "Unsupported step size"); + if constexpr (step_size == 4) { + return (offset ^ 0x4) + step_size * row_stride; + } else { + // step_size % 8 == 0 + return offset + step_size * row_stride; + } + } else { + static_assert(step_size == 4 || step_size % 8 == 0, "Unsupported step size"); + if constexpr (step_size == 4) { + return (offset ^ 0x2) + step_size * row_stride; + } else { + // step_size % 8 == 0 + return offset + step_size * row_stride; + } + } + } + + __device__ __forceinline__ void ldmatrix_m8n8x4(uint32_t offset, uint32_t* R) { + b128_t* smem_ptr = base + offset; + mma::ldmatrix_m8n8x4(R, smem_ptr); + } + + __device__ __forceinline__ void ldmatrix_m8n8x4_left_half(uint32_t offset, uint32_t* R) { + b128_t* smem_ptr = base + offset; + mma::ldmatrix_m8n8x4_left_half(R, smem_ptr); + } + + __device__ __forceinline__ void ldmatrix_m8n8x4_right_half(uint32_t offset, uint32_t* R) { + b128_t* smem_ptr = base + offset; + mma::ldmatrix_m8n8x4_right_half(R, smem_ptr); + } + + __device__ __forceinline__ void stmatrix_m8n8x4(uint32_t offset, uint32_t* R) { + b128_t* smem_ptr = base + offset; + mma::stmatrix_m8n8x4(R, smem_ptr); + } + + __device__ __forceinline__ void ldmatrix_m8n8x4_trans(uint32_t offset, uint32_t* R) { + b128_t* smem_ptr = base + offset; + mma::ldmatrix_m8n8x4_trans(R, smem_ptr); + } + + __device__ __forceinline__ void ldmatrix_m8n8x4_trans_left_half(uint32_t offset, uint32_t* R) { + b128_t* smem_ptr = base + offset; + mma::ldmatrix_m8n8x4_trans_left_half(R, smem_ptr); + } + + __device__ __forceinline__ void ldmatrix_m8n8x4_trans_right_half(uint32_t offset, uint32_t* R) { + b128_t* smem_ptr = base + offset; + mma::ldmatrix_m8n8x4_trans_right_half(R, smem_ptr); + } + + template + __device__ __forceinline__ void load_128b_async(uint32_t offset, const T* gptr, bool predicate) { + b128_t* smem_ptr = base + offset; + cp_async::pred_load_128b( + smem_ptr, reinterpret_cast(gptr), predicate); + } + + template + __device__ __forceinline__ void load_128b_async(uint32_t offset, const T* gptr) { + b128_t* smem_ptr = base + offset; + cp_async::load_128b(smem_ptr, + reinterpret_cast(gptr)); + } + + template + __device__ __forceinline__ void store_128b(uint32_t offset, T* gptr) { + *reinterpret_cast(gptr) = *(base + offset); + } +}; + +} // namespace flashinfer + +#endif // FLASHINFER_PERMUTED_SMEM_CUH_ diff --git a/include/flashinfer/pos_enc.cuh b/include/flashinfer/pos_enc.cuh new file mode 100644 index 0000000000000000000000000000000000000000..083a01aab2eb65351494eba54e8c9d8f57347ad7 --- /dev/null +++ b/include/flashinfer/pos_enc.cuh @@ -0,0 +1,1070 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef FLASHINFER_POS_ENC_CUH_ +#define FLASHINFER_POS_ENC_CUH_ + +#include +#include +#include +#include + +#include "layout.cuh" +#include "math.cuh" +#include "utils.cuh" +#include "vec_dtypes.cuh" + +namespace flashinfer { + +/*! + * \brief An enumeration class that defines different modes for applying RoPE + * (Rotary Positional Embeddings). + */ +enum class PosEncodingMode { + // No rotary positional embeddings + kNone = 0U, + // Apply Llama-style rope. + kRoPELlama = 1U, + // Apply ALiBi bias + kALiBi = 2U +}; + +/*! + * \brief Convert PosEncodingMode to string + * \param pos_encoding_mode A PosEncodingMode value + */ +inline std::string PosEncodingModeToString(const PosEncodingMode& pos_encoding_mode) { + switch (pos_encoding_mode) { + case PosEncodingMode::kNone: + return "None"; + case PosEncodingMode::kRoPELlama: + return "Llama"; + case PosEncodingMode::kALiBi: + return "ALiBi"; + default: + return "Unknown"; + } +} + +__device__ __forceinline__ float get_alibi_slope(uint32_t head_idx, uint32_t num_heads) { + int n = math::ptx_exp2((int)math::ptx_log2(num_heads)); + return head_idx < n ? math::ptx_exp2(-8. * float(head_idx + 1) / float(n)) + : math::ptx_exp2(-4. * float((head_idx + 1 - n) * 2 - 1) / float(n)); +} + +/*! + * \brief Apply RoPE (Rotary Positional Embeddings) to x[0: head_dim], + * return thread-local vector + * \tparam vec_size A template integer indicates the vector size used + * in the kernel + * \tparam bdx A template integer indicates the blockDim.x + * \tparam T A template type indicates the x data type + * \param x A pointer to the start of x data + * \param freq A vector of float indicates the thread-local rope frequency + * \param offset A integer indicates the offset of the position in RoPE + */ +template +__device__ __forceinline__ vec_t vec_apply_llama_rope( + const T* x, const vec_t& freq, int32_t offset, + const uint32_t rotary_dim = vec_size * bdx) { + vec_t permuted_vec, vec; + vec.cast_load(x + threadIdx.x * vec_size); + + if (threadIdx.x * vec_size < rotary_dim) { + permuted_vec.cast_load(x + ((threadIdx.x * vec_size < rotary_dim / 2) + ? threadIdx.x * vec_size + rotary_dim / 2 + : threadIdx.x * vec_size - rotary_dim / 2)); +#pragma unroll + for (uint32_t i = 0; i < vec_size; ++i) { + float embed = float(offset) * freq[i]; + float cos, sin; + __sincosf(embed, &sin, &cos); + vec[i] = + vec[i] * cos + + ((threadIdx.x * vec_size < rotary_dim / 2) ? -permuted_vec[i] : permuted_vec[i]) * sin; + } + } + return vec; +} + +template +__device__ __forceinline__ vec_t vec_apply_llama_rope_cos_sin( + const T* x, const vec_t& cos, const vec_t& sin, + const uint32_t rotary_dim = vec_size * bdx) { + vec_t permuted_vec, vec; + vec.cast_load(x + threadIdx.x * vec_size); + + if (threadIdx.x * vec_size < rotary_dim) { + permuted_vec.cast_load(x + ((threadIdx.x * vec_size < rotary_dim / 2) + ? threadIdx.x * vec_size + rotary_dim / 2 + : threadIdx.x * vec_size - rotary_dim / 2)); +#pragma unroll + for (uint32_t i = 0; i < vec_size; ++i) { + vec[i] = + vec[i] * cos[i] + + ((threadIdx.x * vec_size < rotary_dim / 2) ? -permuted_vec[i] : permuted_vec[i]) * sin[i]; + } + } + return vec; +} + +/*! + * \brief Apply RoPE (Rotary Positional Embeddings) to x[0: head_dim] with interleave, + * return thread-local vector. + * \tparam vec_size A template integer indicates the vector size used + * in the kernel + * \tparam bdx A template integer indicates the blockDim.x + * \tparam T A template type indicates the x data type + * \param x A pointer to the start of x data + * \param freq A vector of float indicates the thread-local rope frequency + * \param offset A integer indicates the offset of the position in RoPE + */ +template +__device__ __forceinline__ vec_t vec_apply_llama_rope_interleave( + const T* x, const vec_t& freq, int32_t offset, + const uint32_t rotary_dim = vec_size * bdx) { + vec_t vec, vec_before; + vec.cast_load(x + threadIdx.x * vec_size); + + if (threadIdx.x * vec_size < rotary_dim) { + vec_before = vec; +#pragma unroll + for (uint32_t i = 0; i < vec_size; ++i) { + float embed = float(offset) * freq[i]; + float cos, sin; + __sincosf(embed, &sin, &cos); + vec[i] = vec[i] * cos + ((i % 2 == 0) ? -vec_before[i ^ 1] : vec_before[i ^ 1]) * sin; + } + } + return vec; +} + +template +__device__ __forceinline__ vec_t vec_apply_llama_rope_cos_sin_interleave( + const T* x, const vec_t& cos, const vec_t& sin, + const uint32_t rotary_dim = vec_size * bdx) { + vec_t vec, vec_before; + vec.cast_load(x + threadIdx.x * vec_size); + + if (threadIdx.x * vec_size < rotary_dim) { + vec_before = vec; +#pragma unroll + for (uint32_t i = 0; i < vec_size; ++i) { + vec[i] = vec[i] * cos[i] + ((i % 2 == 0) ? -vec_before[i ^ 1] : vec_before[i ^ 1]) * sin[i]; + } + } + return vec; +} + +/* +HACK (ByronHsu): in the interleave mode with cos_sin_cache, we actually only use the first half of +cos and sin + +For example, +In the below example, the vec_size is 4 +the computation in the kernel is: + [x1, x2, x3, x4...] * [cos1, cos1, cos2, cos2] + [-x2, x1, -x4, x3...] * [sin1, sin1, sin2, +sin2] the data we loaded are: + - loaded vec = [x1, x2, x3, x4] + - loaded cos = [cos1, cos2, cos3, cos4] + - loaded sin = [sin1, sin2, sin3, sin4] +But only the first half of cos and sin is used in the computation. + +However, we argue the additional overhead is acceptable: + 1. loading additional elements of cos and sin is not adding much overhead. The arithmetic +intensity is the same as non-interleave mode. Each elements of cos and sin is load twice + 2. we don't want two code paths of cos and sin vector for interleave and non-interleave mode. +*/ +template +__device__ __forceinline__ vec_t +vec_apply_llama_rope_cos_sin_interleave_reuse_half(const T* x, const vec_t& cos, + const vec_t& sin, + const uint32_t rotary_dim = vec_size * bdx) { + vec_t vec, vec_before; + vec.cast_load(x + threadIdx.x * vec_size); + + if (threadIdx.x * vec_size < rotary_dim) { + vec_before = vec; +#pragma unroll + for (uint32_t i = 0; i < vec_size; ++i) { + // i / 2 is to get the index of the first half of cos and sin + vec[i] = vec[i] * cos[i / 2] + + ((i % 2 == 0) ? -vec_before[i ^ 1] : vec_before[i ^ 1]) * sin[i / 2]; + } + } + return vec; +} + +template +__global__ void BatchQKApplyRotaryPosIdsCosSinCacheHeadParallelismKernel( + DType* q, DType* k, DType* q_rope, DType* k_rope, float* __restrict__ cos_sin_cache, + IdType* __restrict__ pos_ids, uint32_t nnz, uint32_t num_qo_heads, uint32_t num_kv_heads, + uint32_t rotary_dim, size_t q_stride_n, size_t q_stride_h, size_t k_stride_n, size_t k_stride_h, + size_t q_rope_stride_n, size_t q_rope_stride_h, size_t k_rope_stride_n, + size_t k_rope_stride_h) { + uint32_t bx = blockIdx.x, tx = threadIdx.x, ty = threadIdx.y; + uint32_t by = blockIdx.y; + const uint32_t bdy = blockDim.y; + + vec_t cos, sin; + if (bx * bdy + ty < nnz) { + const uint32_t idx = bx * bdy + ty; + const IdType pos = pos_ids[idx]; + + const int half_rotary_dim = rotary_dim / 2; + + // 1. if interleave: + // - cos = cos_sin_cache[pos_id][tx * vec_size // 2] + // - sin = cos_sin_cache[pos_id][(rot_dim // 2) + tx * vec_size // 2] + // 2. if not interleave + // - cos = cos_cache[pos_id][(tx * vec_size) % (rot_dim // 2)] + // - sin = sin_cache[pos_id][(rot_dim // 2) + (tx * vec_size) % (rot_dim // 2)] + if (tx * vec_size < rotary_dim) { + int sin_offset = rotary_dim / 2; + int vec_idx; + if constexpr (interleave) { + vec_idx = (tx * vec_size) / 2; // Force integer division + } else { + vec_idx = (tx * vec_size) % half_rotary_dim; // Use half_rotary_dim + } + cos.load(cos_sin_cache + (pos * rotary_dim) + vec_idx); + sin.load(cos_sin_cache + (pos * rotary_dim) + (sin_offset + vec_idx)); + } + + if (by < num_qo_heads) { + uint32_t qo_head_idx = by; + DType* q_ptr = q + get_elem_offset_impl(idx, qo_head_idx, 0, q_stride_n, q_stride_h); + DType* q_rope_ptr = + q_rope + get_elem_offset_impl(idx, qo_head_idx, 0, q_rope_stride_n, q_rope_stride_h); + vec_t q_vec; + if constexpr (interleave) { + q_vec = vec_apply_llama_rope_cos_sin_interleave_reuse_half(q_ptr, cos, sin, + rotary_dim); + } else { + q_vec = vec_apply_llama_rope_cos_sin(q_ptr, cos, sin, rotary_dim); + } + q_vec.cast_store(q_rope_ptr + tx * vec_size); + } else { + uint32_t kv_head_idx = by - num_qo_heads; + DType* k_ptr = k + get_elem_offset_impl(idx, kv_head_idx, 0, k_stride_n, k_stride_h); + DType* k_rope_ptr = + k_rope + get_elem_offset_impl(idx, kv_head_idx, 0, k_rope_stride_n, k_rope_stride_h); + vec_t k_vec; + if constexpr (interleave) { + k_vec = vec_apply_llama_rope_cos_sin_interleave_reuse_half(k_ptr, cos, sin, + rotary_dim); + } else { + k_vec = vec_apply_llama_rope_cos_sin(k_ptr, cos, sin, rotary_dim); + } + k_vec.cast_store(k_rope_ptr + tx * vec_size); + } + } +} + +template +__global__ void BatchQKApplyRotaryPosIdsCosSinCacheKernel( + DType* q, DType* k, DType* q_rope, DType* k_rope, float* __restrict__ cos_sin_cache, + IdType* __restrict__ pos_ids, uint32_t nnz, uint32_t num_qo_heads, uint32_t num_kv_heads, + uint32_t rotary_dim, size_t q_stride_n, size_t q_stride_h, size_t k_stride_n, size_t k_stride_h, + size_t q_rope_stride_n, size_t q_rope_stride_h, size_t k_rope_stride_n, + size_t k_rope_stride_h) { + uint32_t bx = blockIdx.x, tx = threadIdx.x, ty = threadIdx.y; + const uint32_t bdy = blockDim.y; + + vec_t cos, sin; + if (bx * bdy + ty < nnz) { + const uint32_t idx = bx * bdy + ty; + const IdType pos = pos_ids[idx]; + const int half_rotary_dim = rotary_dim / 2; + + // 1. if interleave: + // - cos = cos_sin_cache[pos_id][tx * vec_size // 2] + // - sin = cos_sin_cache[pos_id][(rot_dim // 2) + tx * vec_size // 2] + // 2. if not interleave + // - cos = cos_cache[pos_id][(tx * vec_size) % (rot_dim // 2)] + // - sin = sin_cache[pos_id][(rot_dim // 2) + (tx * vec_size) % (rot_dim // 2)] + if (tx * vec_size < rotary_dim) { + int sin_offset = rotary_dim / 2; + int vec_idx; + if constexpr (interleave) { + vec_idx = (tx * vec_size) / 2; // Force integer division + } else { + vec_idx = (tx * vec_size) % half_rotary_dim; // Use half_rotary_dim + } + cos.load(cos_sin_cache + (pos * rotary_dim) + vec_idx); + sin.load(cos_sin_cache + (pos * rotary_dim) + (sin_offset + vec_idx)); + } + + // not to unroll the loop, because num head might be large and might lead to worse performance +#pragma unroll 1 + for (uint32_t qo_head_idx = 0; qo_head_idx < num_qo_heads; ++qo_head_idx) { + DType* q_ptr = q + get_elem_offset_impl(idx, qo_head_idx, 0, q_stride_n, q_stride_h); + DType* q_rope_ptr = + q_rope + get_elem_offset_impl(idx, qo_head_idx, 0, q_rope_stride_n, q_rope_stride_h); + vec_t q_vec; + if constexpr (interleave) { + q_vec = vec_apply_llama_rope_cos_sin_interleave_reuse_half(q_ptr, cos, sin, + rotary_dim); + } else { + q_vec = vec_apply_llama_rope_cos_sin(q_ptr, cos, sin, rotary_dim); + } + q_vec.cast_store(q_rope_ptr + tx * vec_size); + } + +#pragma unroll 1 + for (uint32_t kv_head_idx = 0; kv_head_idx < num_kv_heads; ++kv_head_idx) { + DType* k_ptr = k + get_elem_offset_impl(idx, kv_head_idx, 0, k_stride_n, k_stride_h); + DType* k_rope_ptr = + k_rope + get_elem_offset_impl(idx, kv_head_idx, 0, k_rope_stride_n, k_rope_stride_h); + vec_t k_vec; + if constexpr (interleave) { + k_vec = vec_apply_llama_rope_cos_sin_interleave_reuse_half(k_ptr, cos, sin, + rotary_dim); + } else { + k_vec = vec_apply_llama_rope_cos_sin(k_ptr, cos, sin, rotary_dim); + } + k_vec.cast_store(k_rope_ptr + tx * vec_size); + } + } +} + +template +__global__ void MLARopeQuantizeKernel( + DType* q_rope_in, DType* k_rope_in, DType* q_nope_in, DType* k_nope_in, QuantType* q_rope_out, + QuantType* k_rope_out, QuantType* q_nope_out, QuantType* k_nope_out, + float* __restrict__ cos_sin_cache, IdType* __restrict__ pos_ids, uint32_t nnz, + uint32_t num_heads, size_t q_rope_in_stride_n, size_t q_rope_in_stride_h, + size_t q_nope_in_stride_n, size_t q_nope_in_stride_h, size_t q_rope_out_stride_n, + size_t q_rope_out_stride_h, size_t q_nope_out_stride_n, size_t q_nope_out_stride_h, + size_t k_rope_in_stride, size_t k_nope_in_stride, size_t k_rope_out_stride, + size_t k_nope_out_stride, float quant_scale_q, float quant_scale_kv) { + uint32_t bx = blockIdx.x, tx = threadIdx.x, ty = threadIdx.y; + uint32_t by = blockIdx.y; + uint32_t bdy = blockDim.y; + constexpr uint32_t rotary_dim = 64; + + vec_t cos, sin; + if (bx * bdy + ty < nnz) { + const uint32_t idx = bx * bdy + ty; + const IdType pos = pos_ids[idx]; + + const int half_rotary_dim = rotary_dim / 2; + // 1. if interleave: + // - cos = cos_sin_cache[pos_id][tx * vec_size // 2] + // - sin = cos_sin_cache[pos_id][(rot_dim // 2) + tx * vec_size // 2] + // 2. if not interleave + // - cos = cos_cache[pos_id][(tx * vec_size) % (rot_dim // 2)] + // - sin = sin_cache[pos_id][(rot_dim // 2) + (tx * vec_size) % (rot_dim // 2)] + if (tx * vec_size < rotary_dim) { + int sin_offset = rotary_dim / 2; + int vec_idx; + if constexpr (interleave) { + vec_idx = (tx * vec_size) / 2; // Force integer division + } else { + vec_idx = (tx * vec_size) % half_rotary_dim; // Use half_rotary_dim + } + cos.load(cos_sin_cache + (pos * rotary_dim) + vec_idx); + sin.load(cos_sin_cache + (pos * rotary_dim) + (sin_offset + vec_idx)); + } + + if (by < num_heads) { + // Query RoPE, 64 dim + // allocate (num_heads,) blocks on blockDim.y + uint32_t q_head_idx = by; + DType* q_rope_in_ptr = + q_rope_in + get_elem_offset_impl(idx, q_head_idx, /*elem_idx=*/0, q_rope_in_stride_n, + q_rope_in_stride_h); + QuantType* q_rope_out_ptr = + q_rope_out + get_elem_offset_impl(idx, q_head_idx, /*elem_idx=*/0, q_rope_out_stride_n, + q_rope_out_stride_h); + vec_t q_rope_vec; + if constexpr (interleave) { + q_rope_vec = vec_apply_llama_rope_cos_sin_interleave_reuse_half( + q_rope_in_ptr, cos, sin, rotary_dim); + } else { + q_rope_vec = + vec_apply_llama_rope_cos_sin(q_rope_in_ptr, cos, sin, rotary_dim); + } +#pragma unroll + for (uint32_t i = 0; i < vec_size; ++i) { + q_rope_vec[i] = q_rope_vec[i] * quant_scale_q; + } + q_rope_vec.cast_store(q_rope_out_ptr + tx * vec_size); + } else if (by == num_heads) { + // k/v RoPE, 64 dim + // allocate (1,) blocks on blockDim.y + DType* k_rope_in_ptr = k_rope_in + get_elem_offset_impl(idx, /*head_idx=*/0, /*elem_idx=*/0, + k_rope_in_stride, k_rope_in_stride); + QuantType* k_rope_out_ptr = + k_rope_out + get_elem_offset_impl(idx, /*head_idx=*/0, /*elem_idx=*/0, k_rope_out_stride, + k_rope_out_stride); + vec_t k_rope_vec; + if constexpr (interleave) { + k_rope_vec = vec_apply_llama_rope_cos_sin_interleave_reuse_half( + k_rope_in_ptr, cos, sin, rotary_dim); + } else { + k_rope_vec = + vec_apply_llama_rope_cos_sin(k_rope_in_ptr, cos, sin, rotary_dim); + } +#pragma unroll + for (uint32_t i = 0; i < vec_size; ++i) { + k_rope_vec[i] = k_rope_vec[i] * quant_scale_kv; + } + k_rope_vec.cast_store(k_rope_out_ptr + tx * vec_size); + } else if (by <= num_heads + 8) { + // K/v Non-RoPE part, 512 dim + // allocate (8,) blocks on blockDim.y + uint32_t chunk_idx = (by - num_heads - 1); + DType* k_nope_in_ptr = + k_nope_in + get_elem_offset_impl(idx, /*head_idx=*/0, /*elem_idx=*/64 * chunk_idx, + k_nope_in_stride, k_nope_in_stride); + QuantType* k_nope_out_ptr = + k_nope_out + get_elem_offset_impl(idx, /*head_idx=*/0, /*elem_idx=*/64 * chunk_idx, + k_nope_out_stride, k_nope_out_stride); + vec_t k_nope_vec; + k_nope_vec.cast_load(k_nope_in_ptr + tx * vec_size); +#pragma unroll + for (uint32_t i = 0; i < vec_size; ++i) { + k_nope_vec[i] = k_nope_vec[i] * quant_scale_kv; + } + k_nope_vec.cast_store(k_nope_out_ptr + tx * vec_size); + } else { + // Query Non-RoPE part, 512 dim + // allocate (num_heads * 8,) blocks on blockDim.y + uint32_t q_head_idx = (by - num_heads - 8 - 1) / 8; + uint32_t chunk_idx = (by - num_heads - 8 - 1) % 8; + DType* q_nope_in_ptr = + q_nope_in + get_elem_offset_impl(idx, q_head_idx, /*elem_idx=*/64 * chunk_idx, + q_nope_in_stride_n, q_nope_in_stride_h); + QuantType* q_nope_out_ptr = + q_nope_out + get_elem_offset_impl(idx, q_head_idx, /*elem_idx=*/64 * chunk_idx, + q_nope_out_stride_n, q_nope_out_stride_h); + vec_t q_nope_vec; + q_nope_vec.cast_load(q_nope_in_ptr + tx * vec_size); +#pragma unroll + for (uint32_t i = 0; i < vec_size; ++i) { + q_nope_vec[i] = q_nope_vec[i] * quant_scale_q; + } + q_nope_vec.cast_store(q_nope_out_ptr + tx * vec_size); + } + } +} + +template +__global__ void BatchQKApplyRotaryPosIdsHeadParallelismKernel( + DType* q, DType* k, DType* q_rope, DType* k_rope, IdType* __restrict__ pos_ids, uint32_t nnz, + uint32_t num_qo_heads, uint32_t num_kv_heads, uint32_t rotary_dim, size_t q_stride_n, + size_t q_stride_h, size_t k_stride_n, size_t k_stride_h, size_t q_rope_stride_n, + size_t q_rope_stride_h, size_t k_rope_stride_n, size_t k_rope_stride_h, float smooth_a, + float smooth_b, float rope_rcp_scale, float rope_rcp_theta) { + // NOTE: q and q_rope may be the same ptr, so do k and k_rope + uint32_t bx = blockIdx.x, tx = threadIdx.x, ty = threadIdx.y; + uint32_t by = blockIdx.y; + const uint32_t bdy = blockDim.y; + vec_t freq; + if (tx * vec_size < rotary_dim) { +#pragma unroll + for (uint32_t i = 0; i < vec_size; ++i) { + if constexpr (interleave) { + freq[i] = __powf(rope_rcp_theta, float(2 * ((tx * vec_size + i) / 2)) / float(rotary_dim)); + } else { + freq[i] = __powf(rope_rcp_theta, + float(2 * ((tx * vec_size + i) % (rotary_dim / 2))) / float(rotary_dim)); + } + + float smooth = freq[i] * smooth_a + smooth_b; + smooth = max(0.0f, min(1.0f, smooth)); // clamp to [0, 1] + freq[i] = (1 - smooth) * (freq[i] * rope_rcp_scale) + smooth * freq[i]; + } + } + + vec_t cos, sin; + + if (bx * bdy + ty < nnz) { + const uint32_t idx = bx * bdy + ty; + const IdType pos = pos_ids[idx]; + + if (tx * vec_size < rotary_dim) { +#pragma unroll + for (uint32_t i = 0; i < vec_size; ++i) { + float embed = float(pos) * freq[i]; + __sincosf(embed, &sin[i], &cos[i]); + } + } + + if (by < num_qo_heads) { + uint32_t qo_head_idx = by; + DType* q_ptr = q + get_elem_offset_impl(idx, qo_head_idx, 0, q_stride_n, q_stride_h); + DType* q_rope_ptr = + q_rope + get_elem_offset_impl(idx, qo_head_idx, 0, q_rope_stride_n, q_rope_stride_h); + vec_t q_vec; + if constexpr (interleave) { + q_vec = vec_apply_llama_rope_cos_sin_interleave(q_ptr, cos, sin, rotary_dim); + } else { + q_vec = vec_apply_llama_rope_cos_sin(q_ptr, cos, sin, rotary_dim); + } + q_vec.cast_store(q_rope_ptr + tx * vec_size); + } else { + uint32_t kv_head_idx = by - num_qo_heads; + DType* k_ptr = k + get_elem_offset_impl(idx, kv_head_idx, 0, k_stride_n, k_stride_h); + DType* k_rope_ptr = + k_rope + get_elem_offset_impl(idx, kv_head_idx, 0, k_rope_stride_n, k_rope_stride_h); + vec_t k_vec; + if constexpr (interleave) { + k_vec = vec_apply_llama_rope_cos_sin_interleave(k_ptr, cos, sin, rotary_dim); + } else { + k_vec = vec_apply_llama_rope_cos_sin(k_ptr, cos, sin, rotary_dim); + } + k_vec.cast_store(k_rope_ptr + tx * vec_size); + } + } +} + +template +__global__ void BatchQKApplyRotaryPosIdsKernel( + DType* q, DType* k, DType* q_rope, DType* k_rope, IdType* __restrict__ pos_ids, uint32_t nnz, + uint32_t num_qo_heads, uint32_t num_kv_heads, uint32_t rotary_dim, size_t q_stride_n, + size_t q_stride_h, size_t k_stride_n, size_t k_stride_h, size_t q_rope_stride_n, + size_t q_rope_stride_h, size_t k_rope_stride_n, size_t k_rope_stride_h, float smooth_a, + float smooth_b, float rope_rcp_scale, float rope_rcp_theta) { + // NOTE: q and q_rope may be the same ptr, so do k and k_rope + uint32_t bx = blockIdx.x, tx = threadIdx.x, ty = threadIdx.y; + const uint32_t bdy = blockDim.y; + vec_t freq; + if (tx * vec_size < rotary_dim) { +#pragma unroll + for (uint32_t i = 0; i < vec_size; ++i) { + if constexpr (interleave) { + freq[i] = __powf(rope_rcp_theta, float(2 * ((tx * vec_size + i) / 2)) / float(rotary_dim)); + } else { + freq[i] = __powf(rope_rcp_theta, + float(2 * ((tx * vec_size + i) % (rotary_dim / 2))) / float(rotary_dim)); + } + + float smooth = freq[i] * smooth_a + smooth_b; + smooth = max(0.0f, min(1.0f, smooth)); // clamp to [0, 1] + freq[i] = (1 - smooth) * (freq[i] * rope_rcp_scale) + smooth * freq[i]; + } + } + + vec_t cos, sin; + + if (bx * bdy + ty < nnz) { + const uint32_t idx = bx * bdy + ty; + const IdType pos = pos_ids[idx]; + + if (tx * vec_size < rotary_dim) { +#pragma unroll + for (uint32_t i = 0; i < vec_size; ++i) { + float embed = float(pos) * freq[i]; + __sincosf(embed, &sin[i], &cos[i]); + } + } + +#pragma unroll 1 + for (uint32_t qo_head_idx = 0; qo_head_idx < num_qo_heads; ++qo_head_idx) { + DType* q_ptr = q + get_elem_offset_impl(idx, qo_head_idx, 0, q_stride_n, q_stride_h); + DType* q_rope_ptr = + q_rope + get_elem_offset_impl(idx, qo_head_idx, 0, q_rope_stride_n, q_rope_stride_h); + vec_t q_vec; + if constexpr (interleave) { + q_vec = vec_apply_llama_rope_cos_sin_interleave(q_ptr, cos, sin, rotary_dim); + } else { + q_vec = vec_apply_llama_rope_cos_sin(q_ptr, cos, sin, rotary_dim); + } + q_vec.cast_store(q_rope_ptr + tx * vec_size); + } + +#pragma unroll 1 + for (uint32_t kv_head_idx = 0; kv_head_idx < num_kv_heads; ++kv_head_idx) { + DType* k_ptr = k + get_elem_offset_impl(idx, kv_head_idx, 0, k_stride_n, k_stride_h); + DType* k_rope_ptr = + k_rope + get_elem_offset_impl(idx, kv_head_idx, 0, k_rope_stride_n, k_rope_stride_h); + vec_t k_vec; + if constexpr (interleave) { + k_vec = vec_apply_llama_rope_cos_sin_interleave(k_ptr, cos, sin, rotary_dim); + } else { + k_vec = vec_apply_llama_rope_cos_sin(k_ptr, cos, sin, rotary_dim); + } + k_vec.cast_store(k_rope_ptr + tx * vec_size); + } + } +} + +template +__global__ void BatchQKApplyRotaryKernel( + DType* q, DType* k, DType* q_rope, DType* k_rope, IdType* __restrict__ indptr, + IdType* __restrict__ offsets, uint32_t batch_size, uint32_t num_qo_heads, uint32_t num_kv_heads, + uint32_t rotary_dim, size_t q_stride_n, size_t q_stride_h, size_t k_stride_n, size_t k_stride_h, + size_t q_rope_stride_n, size_t q_rope_stride_h, size_t k_rope_stride_n, size_t k_rope_stride_h, + float smooth_a, float smooth_b, float rope_rcp_scale, float rope_rcp_theta) { + uint32_t bx = blockIdx.x, tx = threadIdx.x, ty = threadIdx.y; + const uint32_t bdy = blockDim.y; + vec_t freq; + if (tx * vec_size < rotary_dim) { +#pragma unroll + for (uint32_t i = 0; i < vec_size; ++i) { + if constexpr (interleave) { + freq[i] = __powf(rope_rcp_theta, float(2 * ((tx * vec_size + i) / 2)) / float(rotary_dim)); + } else { + freq[i] = __powf(rope_rcp_theta, + float(2 * ((tx * vec_size + i) % (rotary_dim / 2))) / float(rotary_dim)); + } + + float smooth = freq[i] * smooth_a + smooth_b; + smooth = max(0.0f, min(1.0f, smooth)); // clamp to [0, 1] + freq[i] = (1 - smooth) * (freq[i] * rope_rcp_scale) + smooth * freq[i]; + } + } + + if (bx < batch_size * num_qo_heads) { + // apply rotary to q + const uint32_t batch_idx = bx / num_qo_heads; + const uint32_t qo_head_idx = bx % num_qo_heads; + const uint32_t seq_len = indptr[batch_idx + 1] - indptr[batch_idx]; + const uint32_t offset = offsets[batch_idx]; +#pragma unroll 2 + for (uint32_t i = 0; i < (seq_len + bdy - 1) / bdy; ++i) { + vec_t q_vec; + if (i * bdy + ty < seq_len) { + DType* q_ptr = q + get_elem_offset_impl(indptr[batch_idx] + i * bdy + ty, qo_head_idx, 0, + q_stride_n, q_stride_h); + DType* q_rope_ptr = + q_rope + get_elem_offset_impl(indptr[batch_idx] + i * bdy + ty, qo_head_idx, 0, + q_rope_stride_n, q_rope_stride_h); + if constexpr (interleave) { + q_vec = vec_apply_llama_rope_interleave(q_ptr, freq, offset + i * bdy + ty, + rotary_dim); + } else { + q_vec = + vec_apply_llama_rope(q_ptr, freq, offset + i * bdy + ty, rotary_dim); + } + q_vec.cast_store(q_rope_ptr + tx * vec_size); + } + } + } else { + // apply rotary to k + uint32_t batch_idx = (bx - batch_size * num_qo_heads) / num_kv_heads; + uint32_t kv_head_idx = (bx - batch_size * num_qo_heads) % num_kv_heads; + const uint32_t seq_len = indptr[batch_idx + 1] - indptr[batch_idx]; + const uint32_t offset = offsets[batch_idx]; +#pragma unroll 2 + for (uint32_t i = 0; i < (seq_len + bdy - 1) / bdy; ++i) { + vec_t k_vec; + if (i * bdy + ty < seq_len) { + DType* k_ptr = k + get_elem_offset_impl(indptr[batch_idx] + i * bdy + ty, kv_head_idx, 0, + k_stride_n, k_stride_h); + DType* k_rope_ptr = + k_rope + get_elem_offset_impl(indptr[batch_idx] + i * bdy + ty, kv_head_idx, 0, + k_rope_stride_n, k_rope_stride_h); + if constexpr (interleave) { + k_vec = vec_apply_llama_rope_interleave(k_ptr, freq, offset + i * bdy + ty, + rotary_dim); + } else { + k_vec = + vec_apply_llama_rope(k_ptr, freq, offset + i * bdy + ty, rotary_dim); + } + k_vec.cast_store(k_rope_ptr + tx * vec_size); + } + } + } +} + +#define DISPATCH_INTERLEAVE(interleave, INTERLEAVE, ...) \ + if (interleave) { \ + const bool INTERLEAVE = true; \ + __VA_ARGS__ \ + } else { \ + const bool INTERLEAVE = false; \ + __VA_ARGS__ \ + } + +template +cudaError_t MLARopeQuantize(DType* q_rope_in, DType* k_rope_in, DType* q_nope_in, DType* k_nope_in, + QuantType* q_rope_out, QuantType* k_rope_out, QuantType* q_nope_out, + QuantType* k_nope_out, float* cos_sin_cache, IdType* pos_ids, + uint32_t nnz, uint32_t num_heads, size_t q_rope_in_stride_n, + size_t q_rope_in_stride_h, size_t q_nope_in_stride_n, + size_t q_nope_in_stride_h, size_t q_rope_out_stride_n, + size_t q_rope_out_stride_h, size_t q_nope_out_stride_n, + size_t q_nope_out_stride_h, size_t k_rope_in_stride, + size_t k_nope_in_stride, size_t k_rope_out_stride, + size_t k_nope_out_stride, float quant_scale_q, float quant_scale_kv, + bool interleave, cudaStream_t stream = nullptr) { + int dev_id = 0; + int num_sms = 0; + FLASHINFER_CUDA_CALL(cudaGetDevice(&dev_id)); + FLASHINFER_CUDA_CALL(cudaDeviceGetAttribute(&num_sms, cudaDevAttrMultiProcessorCount, dev_id)); + + DISPATCH_INTERLEAVE(interleave, INTERLEAVE, { + constexpr uint32_t rotary_dim = 64; + constexpr uint32_t vec_size = 16 / sizeof(DType); + constexpr uint32_t bdx = rotary_dim / vec_size; + uint32_t num_threads = 128U; + uint32_t bdy = num_threads / bdx; + uint32_t nblks_x = (nnz + bdy - 1) / bdy; + + void* args[] = {(void*)&q_rope_in, + (void*)&k_rope_in, + (void*)&q_nope_in, + (void*)&k_nope_in, + (void*)&q_rope_out, + (void*)&k_rope_out, + (void*)&q_nope_out, + (void*)&k_nope_out, + (void*)&cos_sin_cache, + (void*)&pos_ids, + (void*)&nnz, + (void*)&num_heads, + (void*)&q_rope_in_stride_n, + (void*)&q_rope_in_stride_h, + (void*)&q_nope_in_stride_n, + (void*)&q_nope_in_stride_h, + (void*)&q_rope_out_stride_n, + (void*)&q_rope_out_stride_h, + (void*)&q_nope_out_stride_n, + (void*)&q_nope_out_stride_h, + (void*)&k_rope_in_stride, + (void*)&k_nope_in_stride, + (void*)&k_rope_out_stride, + (void*)&k_nope_out_stride, + (void*)&quant_scale_q, + (void*)&quant_scale_kv}; + auto kernel = MLARopeQuantizeKernel; + dim3 nblks(nblks_x, num_heads + 8 + 1 + num_heads * 8); + dim3 nthrs(bdx, bdy); + FLASHINFER_CUDA_CALL(cudaLaunchKernel((void*)kernel, nblks, nthrs, args, 0, stream)); + }); + + return cudaSuccess; +} + +template +cudaError_t BatchQKApplyRotaryPosIdsCosSinCache( + DType* q, DType* k, DType* q_rope, DType* k_rope, float* cos_sin_cache, IdType* pos_ids, + uint32_t nnz, uint32_t num_qo_heads, uint32_t num_kv_heads, uint32_t rotary_dim, + uint32_t head_dim, size_t q_stride_n, size_t q_stride_h, size_t k_stride_n, size_t k_stride_h, + size_t q_rope_stride_n, size_t q_rope_stride_h, size_t k_rope_stride_n, size_t k_rope_stride_h, + bool interleave, cudaStream_t stream = nullptr) { + int dev_id = 0; + int num_sms = 0; + FLASHINFER_CUDA_CALL(cudaGetDevice(&dev_id)); + FLASHINFER_CUDA_CALL(cudaDeviceGetAttribute(&num_sms, cudaDevAttrMultiProcessorCount, dev_id)); + + DISPATCH_INTERLEAVE(interleave, INTERLEAVE, { + DISPATCH_HEAD_DIM(head_dim, HEAD_DIM, { + // operate on 16 Bytes at a time + constexpr uint32_t vec_size = std::max(16 / sizeof(DType), HEAD_DIM / 32); + // how many threads needed per head_dim + constexpr uint32_t bdx = HEAD_DIM / vec_size; + // how many threads needed per block + uint32_t num_threads = std::max(128U, bdx); + // how many tokens can we process in a block + uint32_t bdy = num_threads / bdx; + // how many blocks needed to process all tokens + uint32_t nblks_x = (nnz + bdy - 1) / bdy; + void* args[] = {(void*)&q, + (void*)&k, + (void*)&q_rope, + (void*)&k_rope, + (void*)&cos_sin_cache, + (void*)&pos_ids, + (void*)&nnz, + (void*)&num_qo_heads, + (void*)&num_kv_heads, + (void*)&rotary_dim, + (void*)&q_stride_n, + (void*)&q_stride_h, + (void*)&k_stride_n, + (void*)&k_stride_h, + (void*)&q_rope_stride_n, + (void*)&q_rope_stride_h, + (void*)&k_rope_stride_n, + (void*)&k_rope_stride_h}; + auto kernel_0 = BatchQKApplyRotaryPosIdsCosSinCacheKernel; + + int num_blocks_per_sm_0 = 0; + FLASHINFER_CUDA_CALL(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &num_blocks_per_sm_0, kernel_0, num_threads, /*smem_size=*/0)); + uint32_t num_ctas_0 = num_blocks_per_sm_0 * num_sms; + + if ((nnz + bdy - 1) / bdy >= num_ctas_0) { + dim3 nblks(nblks_x); + dim3 nthrs(bdx, bdy); + FLASHINFER_CUDA_CALL(cudaLaunchKernel((void*)kernel_0, nblks, nthrs, args, 0, stream)); + } else { + dim3 nblks(nblks_x, num_qo_heads + num_kv_heads); + dim3 nthrs(bdx, bdy); + auto kernel_1 = + BatchQKApplyRotaryPosIdsCosSinCacheHeadParallelismKernel; + FLASHINFER_CUDA_CALL(cudaLaunchKernel((void*)kernel_1, nblks, nthrs, args, 0, stream)); + } + }); + }); + + return cudaSuccess; +} + +template +cudaError_t BatchQKApplyRotaryPosIds( + DType* q, DType* k, DType* q_rope, DType* k_rope, IdType* __restrict__ pos_ids, uint32_t nnz, + uint32_t num_qo_heads, uint32_t num_kv_heads, uint32_t rotary_dim, uint32_t head_dim, + size_t q_stride_n, size_t q_stride_h, size_t k_stride_n, size_t k_stride_h, + size_t q_rope_stride_n, size_t q_rope_stride_h, size_t k_rope_stride_n, size_t k_rope_stride_h, + bool interleave, float rope_scale, float rope_theta, cudaStream_t stream = nullptr) { + float rope_rcp_scale = 1.0f / rope_scale; + float rope_rcp_theta = 1.0f / rope_theta; + float smooth_a = 0.f; + float smooth_b = 0.f; + int dev_id = 0; + int num_sms = 0; + FLASHINFER_CUDA_CALL(cudaGetDevice(&dev_id)); + FLASHINFER_CUDA_CALL(cudaDeviceGetAttribute(&num_sms, cudaDevAttrMultiProcessorCount, dev_id)); + + DISPATCH_INTERLEAVE(interleave, INTERLEAVE, { + DISPATCH_HEAD_DIM(head_dim, HEAD_DIM, { + constexpr uint32_t vec_size = std::max(16 / sizeof(DType), HEAD_DIM / 32); + constexpr uint32_t bdx = HEAD_DIM / vec_size; + uint32_t num_threads = std::max(128U, bdx); + uint32_t bdy = num_threads / bdx; + uint32_t nblks_x = (nnz + bdy - 1) / bdy; + + void* args[] = {(void*)&q, + (void*)&k, + (void*)&q_rope, + (void*)&k_rope, + (void*)&pos_ids, + (void*)&nnz, + (void*)&num_qo_heads, + (void*)&num_kv_heads, + (void*)&rotary_dim, + (void*)&q_stride_n, + (void*)&q_stride_h, + (void*)&k_stride_n, + (void*)&k_stride_h, + (void*)&q_rope_stride_n, + (void*)&q_rope_stride_h, + (void*)&k_rope_stride_n, + (void*)&k_rope_stride_h, + (void*)&smooth_a, + (void*)&smooth_b, + (void*)&rope_rcp_scale, + (void*)&rope_rcp_theta}; + auto kernel_0 = + BatchQKApplyRotaryPosIdsKernel; + + int num_blocks_per_sm_0 = 0; + FLASHINFER_CUDA_CALL(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &num_blocks_per_sm_0, kernel_0, num_threads, /*smem_size=*/0)); + uint32_t num_ctas_0 = num_blocks_per_sm_0 * num_sms; + if (nblks_x >= num_ctas_0) { + dim3 nblks(nblks_x); + dim3 nthrs(bdx, bdy); + + FLASHINFER_CUDA_CALL(cudaLaunchKernel((void*)kernel_0, nblks, nthrs, args, 0, stream)); + } else { + dim3 nblks(nblks_x, num_qo_heads + num_kv_heads); + dim3 nthrs(bdx, bdy); + auto kernel_1 = BatchQKApplyRotaryPosIdsHeadParallelismKernel; + + FLASHINFER_CUDA_CALL(cudaLaunchKernel((void*)kernel_1, nblks, nthrs, args, 0, stream)); + } + }); + }); + + return cudaSuccess; +} + +template +cudaError_t BatchQKApplyRotary(DType* q, DType* k, DType* q_rope, DType* k_rope, + IdType* __restrict__ indptr, IdType* __restrict__ offsets, + uint32_t batch_size, uint32_t num_qo_heads, uint32_t num_kv_heads, + uint32_t rotary_dim, uint32_t head_dim, size_t q_stride_n, + size_t q_stride_h, size_t k_stride_n, size_t k_stride_h, + size_t q_rope_stride_n, size_t q_rope_stride_h, + size_t k_rope_stride_n, size_t k_rope_stride_h, bool interleave, + float rope_scale, float rope_theta, cudaStream_t stream = nullptr) { + float rope_rcp_scale = 1.0f / rope_scale; + float rope_rcp_theta = 1.0f / rope_theta; + float smooth_a = 0.f; + float smooth_b = 0.f; + + DISPATCH_INTERLEAVE(interleave, INTERLEAVE, { + DISPATCH_HEAD_DIM(head_dim, HEAD_DIM, { + constexpr uint32_t vec_size = std::max(16 / sizeof(DType), HEAD_DIM / 32); + constexpr uint32_t bdx = HEAD_DIM / vec_size; + uint32_t num_threads = std::max(128U, bdx); + uint32_t bdy = num_threads / bdx; + dim3 nblks(batch_size * (num_qo_heads + num_kv_heads)); + dim3 nthrs(bdx, bdy); + auto kernel = BatchQKApplyRotaryKernel; + void* args[] = {(void*)&q, + (void*)&k, + (void*)&q_rope, + (void*)&k_rope, + (void*)&indptr, + (void*)&offsets, + (void*)&batch_size, + (void*)&num_qo_heads, + (void*)&num_kv_heads, + (void*)&rotary_dim, + (void*)&q_stride_n, + (void*)&q_stride_h, + (void*)&k_stride_n, + (void*)&k_stride_h, + (void*)&q_rope_stride_n, + (void*)&q_rope_stride_h, + (void*)&k_rope_stride_n, + (void*)&k_rope_stride_h, + (void*)&smooth_a, + (void*)&smooth_b, + (void*)&rope_rcp_scale, + (void*)&rope_rcp_theta}; + FLASHINFER_CUDA_CALL(cudaLaunchKernel((void*)kernel, nblks, nthrs, args, 0, stream)); + }); + }); + + return cudaSuccess; +} + +template +cudaError_t BatchQKApplyRotaryInPlace(DType* __restrict__ q, DType* __restrict__ k, + IdType* __restrict__ indptr, IdType* __restrict__ offsets, + uint32_t batch_size, uint32_t num_qo_heads, + uint32_t num_kv_heads, uint32_t rotary_dim, uint32_t head_dim, + size_t q_stride_n, size_t q_stride_h, size_t k_stride_n, + size_t k_stride_h, bool interleave, float rope_scale, + float rope_theta, cudaStream_t stream = nullptr) { + return BatchQKApplyRotary( + q, k, q, k, indptr, offsets, batch_size, num_qo_heads, num_kv_heads, rotary_dim, head_dim, + q_stride_n, q_stride_h, k_stride_n, k_stride_h, q_stride_n, q_stride_h, k_stride_n, + k_stride_h, interleave, rope_scale, rope_theta, stream); +} + +template +cudaError_t BatchQKApplyLlama31Rotary( + DType* q, DType* k, DType* q_rope, DType* k_rope, IdType* __restrict__ indptr, + IdType* __restrict__ offsets, uint32_t batch_size, uint32_t num_qo_heads, uint32_t num_kv_heads, + uint32_t rotary_dim, uint32_t head_dim, size_t q_stride_n, size_t q_stride_h, size_t k_stride_n, + size_t k_stride_h, size_t q_rope_stride_n, size_t q_rope_stride_h, size_t k_rope_stride_n, + size_t k_rope_stride_h, bool interleave, float rope_scale, float rope_theta, + float low_freq_factor, float high_freq_factor, float old_context_length, + cudaStream_t stream = nullptr) { + float rope_rcp_scale = 1.0f / rope_scale; + float rope_rcp_theta = 1.0f / rope_theta; + float smooth_a = old_context_length / (2 * M_PI * high_freq_factor - 2 * M_PI * low_freq_factor); + float smooth_b = -1.0f / (high_freq_factor / low_freq_factor - 1.0f); + + DISPATCH_INTERLEAVE(interleave, INTERLEAVE, { + DISPATCH_HEAD_DIM(head_dim, HEAD_DIM, { + constexpr uint32_t vec_size = std::max(16 / sizeof(DType), HEAD_DIM / 32); + constexpr uint32_t bdx = HEAD_DIM / vec_size; + uint32_t num_threads = std::max(128U, bdx); + uint32_t bdy = num_threads / bdx; + dim3 nblks(batch_size * (num_qo_heads + num_kv_heads)); + dim3 nthrs(bdx, bdy); + auto kernel = BatchQKApplyRotaryKernel; + void* args[] = {(void*)&q, + (void*)&k, + (void*)&q_rope, + (void*)&k_rope, + (void*)&indptr, + (void*)&offsets, + (void*)&batch_size, + (void*)&num_qo_heads, + (void*)&num_kv_heads, + (void*)&rotary_dim, + (void*)&q_stride_n, + (void*)&q_stride_h, + (void*)&k_stride_n, + (void*)&k_stride_h, + (void*)&q_rope_stride_n, + (void*)&q_rope_stride_h, + (void*)&k_rope_stride_n, + (void*)&k_rope_stride_h, + (void*)&smooth_a, + (void*)&smooth_b, + (void*)&rope_rcp_scale, + (void*)&rope_rcp_theta}; + FLASHINFER_CUDA_CALL(cudaLaunchKernel((void*)kernel, nblks, nthrs, args, 0, stream)); + }); + }); + + return cudaSuccess; +} + +template +cudaError_t BatchQKApplyLlama31RotaryPosIds( + DType* q, DType* k, DType* q_rope, DType* k_rope, IdType* pos_ids, uint32_t nnz, + uint32_t num_qo_heads, uint32_t num_kv_heads, uint32_t rotary_dim, uint32_t head_dim, + size_t q_stride_n, size_t q_stride_h, size_t k_stride_n, size_t k_stride_h, + size_t q_rope_stride_n, size_t q_rope_stride_h, size_t k_rope_stride_n, size_t k_rope_stride_h, + bool interleave, float rope_scale, float rope_theta, float low_freq_factor, + float high_freq_factor, float old_context_length, cudaStream_t stream = nullptr) { + float rope_rcp_scale = 1.0f / rope_scale; + float rope_rcp_theta = 1.0f / rope_theta; + float smooth_a = old_context_length / (2 * M_PI * high_freq_factor - 2 * M_PI * low_freq_factor); + float smooth_b = -1.0f / (high_freq_factor / low_freq_factor - 1.0f); + + DISPATCH_INTERLEAVE(interleave, INTERLEAVE, { + DISPATCH_HEAD_DIM(head_dim, HEAD_DIM, { + constexpr uint32_t vec_size = std::max(16 / sizeof(DType), HEAD_DIM / 32); + constexpr uint32_t bdx = HEAD_DIM / vec_size; + uint32_t num_threads = std::max(128U, bdx); + uint32_t bdy = num_threads / bdx; + dim3 nblks((nnz + bdy - 1) / bdy); + dim3 nthrs(bdx, bdy); + auto kernel = + BatchQKApplyRotaryPosIdsKernel; + void* args[] = {(void*)&q, + (void*)&k, + (void*)&q_rope, + (void*)&k_rope, + (void*)&pos_ids, + (void*)&nnz, + (void*)&num_qo_heads, + (void*)&num_kv_heads, + (void*)&rotary_dim, + (void*)&q_stride_n, + (void*)&q_stride_h, + (void*)&k_stride_n, + (void*)&k_stride_h, + (void*)&q_rope_stride_n, + (void*)&q_rope_stride_h, + (void*)&k_rope_stride_n, + (void*)&k_rope_stride_h, + (void*)&smooth_a, + (void*)&smooth_b, + (void*)&rope_rcp_scale, + (void*)&rope_rcp_theta}; + FLASHINFER_CUDA_CALL(cudaLaunchKernel((void*)kernel, nblks, nthrs, args, 0, stream)); + }); + }); + + return cudaSuccess; +} + +} // namespace flashinfer + +#endif // FLASHINFER_POS_ENC_CUH_ diff --git a/include/flashinfer/profiler.cuh b/include/flashinfer/profiler.cuh new file mode 100644 index 0000000000000000000000000000000000000000..9d1cddaba388543d8ed66c5501b1a881407fbd01 --- /dev/null +++ b/include/flashinfer/profiler.cuh @@ -0,0 +1,149 @@ +/* + * Copyright (c) 2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#ifndef FLASHINFER_PROFILER_CUH_ +#define FLASHINFER_PROFILER_CUH_ +#include + +namespace flashinfer { + +__device__ __forceinline__ uint32_t get_block_idx() { + return (blockIdx.z * gridDim.y + blockIdx.y) * gridDim.x + blockIdx.x; +} + +__device__ __forceinline__ uint32_t get_num_blocks() { return gridDim.x * gridDim.y * gridDim.z; } + +__device__ __forceinline__ uint32_t get_thread_idx() { + return (threadIdx.z * blockDim.y + threadIdx.y) * blockDim.x + threadIdx.x; +} + +constexpr uint32_t BLOCK_GROUP_IDX_MASK = 0xFFFFF; +constexpr uint32_t EVENT_IDX_MASK = 0x3FF; +constexpr uint32_t BEGIN_END_MASK = 0x3; + +constexpr uint32_t EVENT_IDX_SHIFT = 2; +constexpr uint32_t BLOCK_GROUP_IDX_SHIFT = 12; +constexpr uint32_t SM_ID_SHIFT = 24; +// Tag layout: +// bits 0-1: event_type (start, end, instant) +// bits 2-11: event_idx (translates to event_names in python profiler) +// bits 12-23: block_id (12 bits) +// bits 24-31: sm_id (8 bits) + +constexpr uint32_t EVENT_BEGIN = 0x0; +constexpr uint32_t EVENT_END = 0x1; +constexpr uint32_t EVENT_INSTANT = 0x2; + +__device__ __forceinline__ uint32_t encode_tag(uint32_t sm_id, uint32_t block_id, + uint32_t event_idx, uint32_t event_type) { + return (sm_id << SM_ID_SHIFT) | (block_id << BLOCK_GROUP_IDX_SHIFT) | + (event_idx << EVENT_IDX_SHIFT) | event_type; +} + +__device__ __forceinline__ uint32_t get_timestamp() { + volatile uint32_t ret; + asm volatile("mov.u32 %0, %globaltimer_lo;" : "=r"(ret)); + return ret; +} + +struct ProfilerEntry { + union { + struct { + uint32_t nblocks; + uint32_t ngroups; + }; + struct { + uint32_t tag; + uint32_t delta_time; + }; + uint64_t raw; + }; +}; + +#ifdef FLASHINFER_ENABLE_PROFILER +#define PROFILER_CLOSURE_PARAMS_DECL \ + ProfilerEntry entry; \ + uint64_t* profiler_write_ptr; \ + uint32_t profiler_write_stride; \ + uint32_t profiler_entry_tag_base; \ + bool profiler_write_thread_predicate; + +#define PROFILER_CLOSURE_FUNC_PARAMS , ProfilerClosure& profiler_closure + +#define PROFILER_FUNC_PARAMS , at::Tensor profiler_buffer +#define PROFILER_PARAMS_DECL uint64_t* profiler_buffer; + +#define PROFILER_INIT(params, smem_storage, closure, group_idx, num_groups, \ + write_thread_predicate) \ + uint32_t _sm_idx; \ + asm volatile("mov.u32 %0, %smid;" : "=r"(_sm_idx)); \ + if (get_block_idx() == 0 && get_thread_idx() == 0) { \ + closure.entry.nblocks = get_num_blocks(); \ + closure.entry.ngroups = num_groups; \ + params.profiler_buffer[0] = closure.entry.raw; \ + } \ + closure.profiler_write_ptr = \ + params.profiler_buffer + 1 + get_block_idx() * num_groups + group_idx; \ + closure.profiler_write_stride = get_num_blocks() * num_groups; \ + closure.profiler_entry_tag_base = encode_tag(_sm_idx, get_block_idx(), 0, 0); \ + closure.profiler_write_thread_predicate = write_thread_predicate; + +#define PROFILER_EVENT_START(closure, event) \ + if (closure.profiler_write_thread_predicate) { \ + closure.entry.tag = \ + closure.profiler_entry_tag_base | ((uint32_t)event << EVENT_IDX_SHIFT) | EVENT_BEGIN; \ + closure.entry.delta_time = get_timestamp(); \ + *closure.profiler_write_ptr = closure.entry.raw; \ + closure.profiler_write_ptr += closure.profiler_write_stride; \ + } \ + __threadfence_block(); + +#define PROFILER_EVENT_END(closure, event) \ + __threadfence_block(); \ + if (closure.profiler_write_thread_predicate) { \ + closure.entry.tag = \ + closure.profiler_entry_tag_base | ((uint32_t)event << EVENT_IDX_SHIFT) | EVENT_END; \ + closure.entry.delta_time = get_timestamp(); \ + *closure.profiler_write_ptr = closure.entry.raw; \ + closure.profiler_write_ptr += closure.profiler_write_stride; \ + } + +#define PROFILER_EVENT_INSTANT(closure, event) \ + __threadfence_block(); \ + if (closure.profiler_write_thread_predicate) { \ + closure.entry.tag = \ + closure.profiler_entry_tag_base | ((uint32_t)event << EVENT_IDX_SHIFT) | EVENT_INSTANT; \ + closure.entry.delta_time = get_timestamp(); \ + *closure.profiler_write_ptr = closure.entry.raw; \ + } \ + __threadfence_block(); + +#else + +#define PROFILER_CLOSURE_PARAMS_DECL +#define PROFILER_CLOSURE_FUNC_PARAMS +#define PROFILER_FUNC_PARAMS +#define PROFILER_PARAMS_DECL +#define PROFILER_INIT(params, smem_storage, closure, group_idx, num_groups, write_thread_predicate) +#define PROFILER_EVENT_START(closure, event) +#define PROFILER_EVENT_END(closure, event) +#define PROFILER_EVENT_INSTANT(closure, event) + +#endif + +} // namespace flashinfer + +#endif // FLASHINFER_PROFILER_CUH_ diff --git a/include/flashinfer/quantization.cuh b/include/flashinfer/quantization.cuh new file mode 100644 index 0000000000000000000000000000000000000000..5043b7e930b78f7fc0fcd76d8544b9c62a8c6abf --- /dev/null +++ b/include/flashinfer/quantization.cuh @@ -0,0 +1,120 @@ +/* + * Copyright (c) 2024 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef FLASHINFER_QUANTIZATION_CUH_ +#define FLASHINFER_QUANTIZATION_CUH_ +#include +#include + +#include + +#include "utils.cuh" + +namespace flashinfer { +namespace quantization { + +enum class BitOrder { kBig = 0U, kLittle = 1U }; + +#define DISPATCH_BITORDER(bitorder, BITORDER, ...) \ + if (bitorder == BitOrder::kBig) { \ + constexpr BitOrder BITORDER = BitOrder::kBig; \ + __VA_ARGS__ \ + } else { \ + constexpr BitOrder BITORDER = BitOrder::kLittle; \ + __VA_ARGS__ \ + } + +template +__global__ void PackBitsKernel(bool* input, uint8_t* output, int64_t num_elements) { + int64_t start_offset = static_cast(blockIdx.x) * blockDim.x * 8, tx = threadIdx.x; + uint8_t ret = 0; + bool input_vec[8]; + typedef cub::BlockLoad BlockLoad; + __shared__ typename BlockLoad::TempStorage temp_storage; + + // This fix the INT32_T overflow issue, which is possible in DiT video models + // where the kv_len could be 128K. + // ref: + // https://github.com/NVIDIA/cub/blob/0fc3c3701632a4be906765b73be20a9ad0da603d/cub/block/block_load.cuh#L711C13-L711C100 + int block_items_end = + (num_elements - start_offset > INT32_MAX) ? INT32_MAX : num_elements - start_offset; + BlockLoad(temp_storage).Load(input + start_offset, input_vec, block_items_end, /*default=*/0); + + if constexpr (BITORDER == BitOrder::kBig) { + ret = (input_vec[0] << 7) | (input_vec[1] << 6) | (input_vec[2] << 5) | (input_vec[3] << 4) | + (input_vec[4] << 3) | (input_vec[5] << 2) | (input_vec[6] << 1) | input_vec[7]; + } else { + ret = (input_vec[7] << 7) | (input_vec[6] << 6) | (input_vec[5] << 5) | (input_vec[4] << 4) | + (input_vec[3] << 3) | (input_vec[2] << 2) | (input_vec[1] << 1) | input_vec[0]; + } + if (start_offset + tx * 8 < num_elements) output[start_offset / 8 + tx] = ret; +} + +template +__global__ void SegmentPackBitsKernel(bool* input, uint8_t* output, IdType* input_indptr, + IdType* output_indptr) { + int64_t bx = blockIdx.x, tx = threadIdx.x; + bool input_vec[8]; + typedef cub::BlockLoad BlockLoad; + __shared__ typename BlockLoad::TempStorage temp_storage; + int64_t num_elements = input_indptr[bx + 1] - input_indptr[bx]; + for (uint32_t start_offset = 0; start_offset < num_elements; start_offset += 8 * blockDim.x) { + uint8_t ret = 0; + BlockLoad(temp_storage) + .Load(input + input_indptr[bx] + start_offset, input_vec, num_elements - start_offset, + /*default=*/0); + + if constexpr (BITORDER == BitOrder::kBig) { + ret = (input_vec[0] << 7) | (input_vec[1] << 6) | (input_vec[2] << 5) | (input_vec[3] << 4) | + (input_vec[4] << 3) | (input_vec[5] << 2) | (input_vec[6] << 1) | input_vec[7]; + } else { + ret = (input_vec[7] << 7) | (input_vec[6] << 6) | (input_vec[5] << 5) | (input_vec[4] << 4) | + (input_vec[3] << 3) | (input_vec[2] << 2) | (input_vec[1] << 1) | input_vec[0]; + } + if (start_offset + tx * 8 < num_elements) + output[output_indptr[bx] + start_offset / 8 + tx] = ret; + } +} + +cudaError_t PackBits(bool* input, uint8_t* output, int64_t num_elements, BitOrder bitorder, + cudaStream_t stream) { + DISPATCH_BITORDER(bitorder, BITORDER, { + auto kernel = PackBitsKernel; + const dim3 nthrs(256); + const dim3 nblks(ceil_div(num_elements, nthrs.x * 8)); + void* args[] = {&input, &output, &num_elements}; + FLASHINFER_CUDA_CALL(cudaLaunchKernel((void*)kernel, nblks, nthrs, args, 0, stream)); + }); + return cudaSuccess; +} + +template +cudaError_t SegmentPackBits(bool* input, uint8_t* output, IdType* input_indptr, + IdType* output_indptr, uint32_t batch_size, BitOrder bitorder, + cudaStream_t stream) { + DISPATCH_BITORDER(bitorder, BITORDER, { + auto kernel = SegmentPackBitsKernel; + const dim3 nthrs(256); + const dim3 nblks(batch_size); + void* args[] = {&input, &output, &input_indptr, &output_indptr}; + FLASHINFER_CUDA_CALL(cudaLaunchKernel((void*)kernel, nblks, nthrs, args, 0, stream)); + }); + return cudaSuccess; +} + +} // namespace quantization +} // namespace flashinfer + +#endif // FLASHINFER_QUANTIZATION_CUH_ diff --git a/include/flashinfer/sampling.cuh b/include/flashinfer/sampling.cuh new file mode 100644 index 0000000000000000000000000000000000000000..7b0002bf33eda4005119894c26bd5df2ef6cf7ae --- /dev/null +++ b/include/flashinfer/sampling.cuh @@ -0,0 +1,2156 @@ +/* + * Copyright (c) 2024 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef FLASHINFER_SAMPLING_CUH_ +#define FLASHINFER_SAMPLING_CUH_ + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "allocator.h" +#include "math.cuh" +#include "utils.cuh" +#include "vec_dtypes.cuh" + +namespace flashinfer { + +namespace sampling { + +using namespace cub; + +#define DISPATCH_DETERMINISTIC(deterministic, DETERMINISTIC, ...) \ + if (deterministic) { \ + constexpr bool DETERMINISTIC = true; \ + __VA_ARGS__ \ + } else { \ + constexpr bool DETERMINISTIC = false; \ + __VA_ARGS__ \ + } + +#define DISPATCH_COMPUTE_CAP_NUM_THREADS(compute_capacity, BLOCK_THREADS, ...) \ + if (compute_capacity.first >= 8) { \ + constexpr uint32_t BLOCK_THREADS = 1024; \ + __VA_ARGS__ \ + } else { \ + constexpr uint32_t BLOCK_THREADS = 512; \ + __VA_ARGS__ \ + } + +#define DISPATCH_SOFTMAX_CACHE_INPUT(cache_input, CACHE_INPUT, ...) \ + if (cache_input) { \ + constexpr bool CACHE_INPUT = true; \ + __VA_ARGS__ \ + } else { \ + constexpr bool CACHE_INPUT = false; \ + __VA_ARGS__ \ + } + +constexpr BlockScanAlgorithm SCAN_ALGO = BLOCK_SCAN_WARP_SCANS; +constexpr BlockReduceAlgorithm REDUCE_ALGO = BLOCK_REDUCE_WARP_REDUCTIONS; + +#if (__CUDACC_VER_MAJOR__ * 10000 + __CUDACC_VER_MINOR__ * 100 >= 120100) +#define FLASHINFER_CUB_SUBTRACTLEFT_DEFINED +#endif + +template +struct ValueCount { + T value; + int count; + + __device__ ValueCount operator+(const ValueCount& other) const { + return {value + other.value, count + other.count}; + } + __device__ ValueCount& operator+=(const ValueCount& other) { + value += other.value; + count += other.count; + return *this; + } +}; + +struct BoolDiffOp { + __device__ __forceinline__ bool operator()(const bool& lhs, const bool& rhs) const { + return lhs != rhs; + } +}; + +struct Float2SoftmaxReduceOp { + __device__ __forceinline__ float2 operator()(const float2& a, const float2& b) const { + if (isinf(a.x)) return b; + if (isinf(b.x)) return a; + + float new_max = max(a.x, b.x); + float new_denom = a.y * __expf(a.x - new_max) + b.y * __expf(b.x - new_max); + return make_float2(new_max, new_denom); + } +}; + +template +struct SamplingTempStorage { + union { + float deterministic_scan[BLOCK_THREADS / 32]; + typename BlockScan::TempStorage scan; + typename BlockReduce::TempStorage reduce; + typename BlockReduce::TempStorage reduce_int; + typename BlockReduce, BLOCK_THREADS, REDUCE_ALGORITHM>::TempStorage + reduce_value_count; + typename BlockAdjacentDifference::TempStorage adj_diff; + } block_prim; + struct { + int32_t sampled_id; + int32_t last_valid_id; + float max_val; + union { + float value; + ValueCount pair; + } block_aggregate; + }; +}; + +template +struct OnlineSoftmaxTempStorage { + union { + typename cub::BlockReduce::TempStorage reduce; + typename cub::BlockReduce::TempStorage reduce_pair; + } block_prim; + + struct { + float max_val; + float denominator; + } shared_state; +}; + +struct PartialSoftmaxResult { + float max_val; + float denominator; +}; + +/*! + * \brief Deterministic inclusive scan implementation, use Belloch scan algorithm. + * \note This implementation is slower than the cub::BlockScan, but it is deterministic. + */ +template +__device__ __forceinline__ void DeterministicInclusiveSum( + const float* in_data, float* out_data, + SamplingTempStorage* temp_storage) { + float* smem_prefix_sum = temp_storage->block_prim.deterministic_scan; + float thread_data[VEC_SIZE]; + float thread_sum = 0; +#pragma unroll + for (uint32_t i = 0; i < VEC_SIZE; ++i) { + thread_sum += in_data[i]; + thread_data[i] = thread_sum; + } + + float thread_exclusive_prefix_sum = thread_sum; + +#pragma unroll + for (uint32_t offset = 1; offset < 32; offset *= 2) { + float tmp = __shfl_up_sync(0xffffffff, thread_exclusive_prefix_sum, offset); + if ((threadIdx.x + 1) % (offset * 2) == 0) { + thread_exclusive_prefix_sum += tmp; + } + } + + float warp_sum = __shfl_sync(0xffffffff, thread_exclusive_prefix_sum, threadIdx.x | 0xffffffff); + if (threadIdx.x % 32 == 31) { + thread_exclusive_prefix_sum = 0; + } + +#pragma unroll + for (uint32_t offset = 16; offset >= 1; offset /= 2) { + float tmp = __shfl_xor_sync(0xffffffff, thread_exclusive_prefix_sum, offset); + if ((threadIdx.x + 1) % (offset * 2) == 0) { + thread_exclusive_prefix_sum = tmp + thread_exclusive_prefix_sum; + } + if ((threadIdx.x + 1) % (offset * 2) == offset) { + thread_exclusive_prefix_sum = tmp; + } + } + + smem_prefix_sum[threadIdx.x / 32] = warp_sum; + __syncthreads(); + + if (threadIdx.x < 32) { + float warp_exclusive_prefix_sum = + (threadIdx.x < BLOCK_THREADS / 32) ? smem_prefix_sum[threadIdx.x] : 0; + +#pragma unroll + for (uint32_t offset = 1; offset < 32; offset *= 2) { + float tmp = __shfl_up_sync(0xffffffff, warp_exclusive_prefix_sum, offset); + if ((threadIdx.x + 1) % (offset * 2) == 0) { + warp_exclusive_prefix_sum += tmp; + } + } + + if (threadIdx.x % 32 == 31) { + warp_exclusive_prefix_sum = 0; + } + +#pragma unroll + for (uint32_t offset = 16; offset >= 1; offset /= 2) { + float tmp = __shfl_xor_sync(0xffffffff, warp_exclusive_prefix_sum, offset); + if ((threadIdx.x + 1) % (offset * 2) == 0) { + warp_exclusive_prefix_sum = tmp + warp_exclusive_prefix_sum; + } + if ((threadIdx.x + 1) % (offset * 2) == offset) { + warp_exclusive_prefix_sum = tmp; + } + } + if (threadIdx.x < BLOCK_THREADS / 32) { + smem_prefix_sum[threadIdx.x] = warp_exclusive_prefix_sum; + } + } + __syncthreads(); + +#pragma unroll + for (uint32_t i = 0; i < VEC_SIZE; ++i) { + out_data[i] = smem_prefix_sum[threadIdx.x / 32] + thread_exclusive_prefix_sum + thread_data[i]; + } +} + +template +__device__ __forceinline__ std::tuple GetMinMaxValue(float* in_data, uint32_t row_idx, + uint32_t d, + TempStorage& temp_storage) { + const uint32_t tx = threadIdx.x; + vec_t in_data_vec; + float max_val = -cuda::std::numeric_limits::infinity(), + min_val = cuda::std::numeric_limits::infinity(); + for (uint32_t i = 0; i < ceil_div(d, BLOCK_THREADS * VEC_SIZE); ++i) { + in_data_vec.fill(0); + if ((i * BLOCK_THREADS + tx) * VEC_SIZE < d) { + in_data_vec.cast_load(in_data + row_idx * d + i * BLOCK_THREADS * VEC_SIZE + tx * VEC_SIZE); + } + float in_data_[VEC_SIZE]; +#pragma unroll + for (uint32_t j = 0; j < VEC_SIZE; ++j) { + in_data_[j] = in_data_vec[j]; + } + max_val = max( + max_val, BlockReduce(temp_storage.block_prim.reduce) + .Reduce(in_data_, cub::Max())); + __syncthreads(); + min_val = min( + min_val, BlockReduce(temp_storage.block_prim.reduce) + .Reduce(in_data_, cub::Min())); + __syncthreads(); + } + if (tx == 0) { + temp_storage.max_val = max_val; + temp_storage.min_val = min_val; + } + __syncthreads(); + max_val = temp_storage.max_val; + min_val = temp_storage.min_val; + + return std::make_tuple(min_val, max_val); +} + +template +__device__ __forceinline__ float GetMaxValue(float* in_data, uint32_t row_idx, uint32_t d, + TempStorage& temp_storage) { + const uint32_t tx = threadIdx.x; + vec_t in_data_vec; + + float max_val = 0; + for (uint32_t i = 0; i < ceil_div(d, BLOCK_THREADS * VEC_SIZE); ++i) { + in_data_vec.fill(0); + if ((i * BLOCK_THREADS + tx) * VEC_SIZE < d) { + in_data_vec.cast_load(in_data + row_idx * d + (i * BLOCK_THREADS + tx) * VEC_SIZE); + } + float in_data_[VEC_SIZE]; +#pragma unroll + for (uint32_t j = 0; j < VEC_SIZE; ++j) { + in_data_[j] = in_data_vec[j]; + } + max_val = max( + max_val, BlockReduce(temp_storage.block_prim.reduce) + .Reduce(in_data_, cub::Max())); + __syncthreads(); + } + if (tx == 0) { + temp_storage.max_val = max_val; + } + __syncthreads(); + return temp_storage.max_val; +} + +template +__global__ void OnlineSoftmaxFusedKernel(DType* logits, DType* output, DType* temperature_arr, + DType temperature_val, uint32_t d) { + const uint32_t bx = blockIdx.x, tx = threadIdx.x; + float temperature = temperature_arr == nullptr ? temperature_val : temperature_arr[bx]; + const float inv_temp = (temperature == 0.f) ? 0.f : 1.f / temperature; + + using TempStorage = OnlineSoftmaxTempStorage; + extern __shared__ __align__(alignof(TempStorage)) uint8_t smem[]; + auto& temp_storage = reinterpret_cast(smem); + + DType* smem_vec_base = nullptr; + if constexpr (CACHE_INPUT) { + constexpr size_t vec_alignment = alignof(vec_t); + size_t aligned_offset = round_up(sizeof(TempStorage), vec_alignment); + smem_vec_base = reinterpret_cast(smem + aligned_offset); + } + + vec_t logits_vec; + + float running_max = -cuda::std::numeric_limits::infinity(); + float running_denominator = 0.0f; + +#if (__CUDACC_VER_MAJOR__ >= 12 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.wait;"); +#endif + + // Pass 1: Compute running max and denominator +#pragma unroll 2 + for (uint32_t i = 0; i < ceil_div(d, BLOCK_THREADS * VEC_SIZE); ++i) { + logits_vec.fill(-cuda::std::numeric_limits::infinity()); + if ((i * BLOCK_THREADS + tx) * VEC_SIZE < d) { + logits_vec.cast_load(logits + bx * d + (i * BLOCK_THREADS + tx) * VEC_SIZE); + +#pragma unroll + for (uint32_t j = 0; j < VEC_SIZE; ++j) { + logits_vec[j] *= inv_temp; + } + + if constexpr (CACHE_INPUT) { + logits_vec.store(smem_vec_base + (i * BLOCK_THREADS + tx) * VEC_SIZE); + } + } + + float thread_max = -cuda::std::numeric_limits::infinity(); +#pragma unroll + for (uint32_t j = 0; j < VEC_SIZE; ++j) { + thread_max = max(thread_max, logits_vec[j]); + } + float block_max = cub::BlockReduce(temp_storage.block_prim.reduce) + .Reduce(thread_max, cub::Max()); + + if (tx == 0) { + temp_storage.shared_state.max_val = block_max; + } + __syncthreads(); + block_max = temp_storage.shared_state.max_val; + + // if block_max is -inf, then this block contains all -inf values, so we can skip updating + if (!isinf(block_max)) { + float thread_sum = 0.0f; +#pragma unroll + for (uint32_t j = 0; j < VEC_SIZE; ++j) { + thread_sum += __expf(logits_vec[j] - block_max); + } + + float block_sum = + cub::BlockReduce(temp_storage.block_prim.reduce).Sum(thread_sum); + __syncthreads(); + + if (tx == 0) { + float new_max = max(running_max, block_max); + running_denominator = running_denominator * __expf(running_max - new_max) + + block_sum * __expf(block_max - new_max); + running_max = new_max; + + temp_storage.shared_state.max_val = running_max; + temp_storage.shared_state.denominator = running_denominator; + } + __syncthreads(); + running_max = temp_storage.shared_state.max_val; + running_denominator = temp_storage.shared_state.denominator; + } + } + + const float final_max = running_max; + const float inv_denominator = 1.0f / running_denominator; + + __syncthreads(); + + // Pass 2: Normalize in place + vec_t prob_vec; + for (uint32_t i = 0; i < ceil_div(d, BLOCK_THREADS * VEC_SIZE); ++i) { + if constexpr (CACHE_INPUT) { + logits_vec.load(smem_vec_base + (i * BLOCK_THREADS + tx) * VEC_SIZE); + } else { + if ((i * BLOCK_THREADS + tx) * VEC_SIZE < d) { + logits_vec.cast_load(logits + bx * d + (i * BLOCK_THREADS + tx) * VEC_SIZE); + +#pragma unroll + for (uint32_t j = 0; j < VEC_SIZE; ++j) { + logits_vec[j] *= inv_temp; + } + } + } + +#pragma unroll + for (uint32_t j = 0; j < VEC_SIZE; ++j) { + float p = __expf(static_cast(logits_vec[j]) - final_max) * inv_denominator; + prob_vec[j] = static_cast(p); + } + + if ((i * BLOCK_THREADS + tx) * VEC_SIZE < d) { + prob_vec.cast_store(output + bx * d + (i * BLOCK_THREADS + tx) * VEC_SIZE); + } + } +#if (__CUDACC_VER_MAJOR__ >= 12 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.launch_dependents;"); +#endif +} + +template +__global__ void OnlineSoftmaxMapKernel(DType* logits, PartialSoftmaxResult* partial_results, + DType* temperature_arr, float temperature_val, uint32_t d, + uint32_t num_slices) { + const uint32_t bx = blockIdx.x; + const uint32_t by = blockIdx.y; // slice index + const uint32_t tx = threadIdx.x; + float temperature = temperature_arr == nullptr ? temperature_val : temperature_arr[bx]; + const float inv_temp = (temperature == 0.f) ? 0.f : 1.f / temperature; + + const uint32_t vec_alignment_elems = alignof(vec_t) / sizeof(DType); + const uint32_t slice_stride = round_up(ceil_div(d, num_slices), vec_alignment_elems); + const uint32_t slice_start = by * slice_stride; + const uint32_t slice_size = min((by + 1) * slice_stride, d) - slice_start; + + if (slice_start >= d) return; + + using TempStorage = OnlineSoftmaxTempStorage; + extern __shared__ __align__(alignof(TempStorage)) uint8_t smem[]; + auto& temp_storage = reinterpret_cast(smem); + + vec_t logits_vec; + float running_max = -cuda::std::numeric_limits::infinity(); + float running_denominator = 0.0f; + +#if (__CUDACC_VER_MAJOR__ >= 12 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.wait;"); +#endif + +#pragma unroll 2 + for (uint32_t i = 0; i < ceil_div(slice_size, BLOCK_THREADS * VEC_SIZE); ++i) { + logits_vec.fill(-cuda::std::numeric_limits::infinity()); + + if ((i * BLOCK_THREADS + tx) * VEC_SIZE < slice_size) { + logits_vec.cast_load(logits + bx * d + slice_start + (i * BLOCK_THREADS + tx) * VEC_SIZE); + } + + float thread_max = -cuda::std::numeric_limits::infinity(); +#pragma unroll + for (uint32_t j = 0; j < VEC_SIZE; ++j) { + logits_vec[j] *= inv_temp; + thread_max = max(thread_max, logits_vec[j]); + } + + float block_max = cub::BlockReduce(temp_storage.block_prim.reduce) + .Reduce(thread_max, cub::Max()); + + if (tx == 0) { + temp_storage.shared_state.max_val = block_max; + } + __syncthreads(); + block_max = temp_storage.shared_state.max_val; + + // if block_max is -inf, then this block contains all -inf values, so we can skip updating + if (!isinf(block_max)) { + float thread_sum = 0.0f; +#pragma unroll + for (uint32_t j = 0; j < VEC_SIZE; ++j) { + thread_sum += __expf(logits_vec[j] - block_max); + } + + float block_sum = + cub::BlockReduce(temp_storage.block_prim.reduce).Sum(thread_sum); + __syncthreads(); + + if (tx == 0) { + float new_max = max(running_max, block_max); + running_denominator = running_denominator * __expf(running_max - new_max) + + block_sum * __expf(block_max - new_max); + running_max = new_max; + + temp_storage.shared_state.max_val = running_max; + temp_storage.shared_state.denominator = running_denominator; + } + __syncthreads(); + running_max = temp_storage.shared_state.max_val; + running_denominator = temp_storage.shared_state.denominator; + } + } + + if (tx == 0) { + partial_results[bx * num_slices + by] = {running_max, running_denominator}; + } +#if (__CUDACC_VER_MAJOR__ >= 12 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.launch_dependents;"); +#endif +} + +template +__global__ void OnlineSoftmaxReduceKernel(DType* logits, DType* output, + PartialSoftmaxResult* partial_results, + DType* temperature_arr, float temperature_val, uint32_t d, + uint32_t num_slices) { + const uint32_t bx = blockIdx.x; + const uint32_t tx = threadIdx.x; + float temperature = temperature_arr == nullptr ? temperature_val : temperature_arr[bx]; + const float inv_temp = (temperature == 0.f) ? 0.f : 1.f / temperature; + + // Reduce slice results + using TempStorage = OnlineSoftmaxTempStorage; + extern __shared__ __align__(alignof(TempStorage)) uint8_t smem[]; + auto& temp_storage = reinterpret_cast(smem); + + const Float2SoftmaxReduceOp reduce_op; + + float2 thread_aggregate = make_float2(-cuda::std::numeric_limits::infinity(), 0.0f); + +#if (__CUDACC_VER_MAJOR__ >= 12 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.wait;"); +#endif + + for (uint32_t i = tx; i < num_slices; i += BLOCK_THREADS) { + PartialSoftmaxResult partial = partial_results[bx * num_slices + i]; + float2 partial_pair = make_float2(partial.max_val, partial.denominator); + thread_aggregate = reduce_op(thread_aggregate, partial_pair); + } + + float2 block_result = cub::BlockReduce(temp_storage.block_prim.reduce_pair) + .Reduce(thread_aggregate, reduce_op); + + if (tx == 0) { + temp_storage.shared_state.max_val = block_result.x; + temp_storage.shared_state.denominator = block_result.y; + } + __syncthreads(); + + block_result = + make_float2(temp_storage.shared_state.max_val, temp_storage.shared_state.denominator); + + const float final_max = temp_storage.shared_state.max_val; + const float inv_denominator = 1.0f / temp_storage.shared_state.denominator; + + // Apply normalization + vec_t logits_vec; + vec_t prob_vec; + + for (uint32_t i = 0; i < ceil_div(d, BLOCK_THREADS * VEC_SIZE); ++i) { + logits_vec.fill(-cuda::std::numeric_limits::infinity()); + + if ((i * BLOCK_THREADS + tx) * VEC_SIZE < d) { + logits_vec.cast_load(logits + bx * d + (i * BLOCK_THREADS + tx) * VEC_SIZE); + } + +#pragma unroll + for (uint32_t j = 0; j < VEC_SIZE; ++j) { + logits_vec[j] *= inv_temp; + float p = __expf(static_cast(logits_vec[j]) - final_max) * inv_denominator; + prob_vec[j] = static_cast(p); + } + + if ((i * BLOCK_THREADS + tx) * VEC_SIZE < d) { + prob_vec.cast_store(output + bx * d + (i * BLOCK_THREADS + tx) * VEC_SIZE); + } + } +#if (__CUDACC_VER_MAJOR__ >= 12 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.launch_dependents;"); +#endif +} + +template +__device__ __forceinline__ void DeviceSamplingFromProb( + uint32_t i, uint32_t d, Predicate pred, float u, vec_t prob_vec, + float& aggregate, + SamplingTempStorage* temp_storage) { + const uint32_t tx = threadIdx.x; + float prob_greater_than_threshold[VEC_SIZE]; + float inclusive_cdf[VEC_SIZE]; + bool greater_than_u[VEC_SIZE], valid[VEC_SIZE]; +#pragma unroll + for (uint32_t j = 0; j < VEC_SIZE; ++j) { + prob_greater_than_threshold[j] = pred(prob_vec[j]) ? prob_vec[j] : 0; + valid[j] = pred(prob_vec[j]) && (i * BLOCK_THREADS + tx) * VEC_SIZE + j < d; + } + float aggregate_local = + BlockReduce(temp_storage->block_prim.reduce) + .Sum(prob_greater_than_threshold); + if (tx == 0) { + temp_storage->block_aggregate.value = aggregate_local; + } + __syncthreads(); + aggregate_local = temp_storage->block_aggregate.value; + + if (aggregate + aggregate_local > u) { + if constexpr (DETERMINISTIC) { + DeterministicInclusiveSum( + prob_greater_than_threshold, inclusive_cdf, temp_storage); + } else { + BlockScan(temp_storage->block_prim.scan) + .InclusiveSum(prob_greater_than_threshold, inclusive_cdf); + + __syncthreads(); + } + +#pragma unroll + for (uint32_t j = 0; j < VEC_SIZE; ++j) { + greater_than_u[j] = (inclusive_cdf[j] + aggregate > u) && valid[j]; + } + + bool greater_than_u_diff[VEC_SIZE]; +#ifdef FLASHINFER_CUB_SUBTRACTLEFT_DEFINED + BlockAdjacentDifference(temp_storage->block_prim.adj_diff) + .SubtractLeft(greater_than_u, greater_than_u_diff, BoolDiffOp()); +#else + BlockAdjacentDifference(temp_storage->block_prim.adj_diff) + .FlagHeads(greater_than_u_diff, greater_than_u, BoolDiffOp(), 0); +#endif + __syncthreads(); + +#pragma unroll + for (uint32_t j = 0; j < VEC_SIZE; ++j) { + if (greater_than_u_diff[j]) { + atomicMin(&(temp_storage->sampled_id), (i * BLOCK_THREADS + tx) * VEC_SIZE + j); + } + } + __syncthreads(); + } + + // update the last valid index + int valid_index[VEC_SIZE]; +#pragma unroll + for (uint32_t j = 0; j < VEC_SIZE; ++j) { + if (valid[j]) { + valid_index[j] = (i * BLOCK_THREADS + tx) * VEC_SIZE + j; + } else { + valid_index[j] = -1; + } + } + int max_valid_index = + BlockReduce(temp_storage->block_prim.reduce_int) + .Reduce(valid_index, cub::Max()); + if (tx == 0 && max_valid_index != -1) { + temp_storage->last_valid_id = max_valid_index; + } + __syncthreads(); + aggregate += aggregate_local; +} + +template +struct DataAndIndex { + DType data; + IdType index; + + __device__ DataAndIndex operator+(const DataAndIndex& other) const { + if (data > other.data) { + return {data, index}; + } else { + return {other.data, other.index}; + } + } + __device__ DataAndIndex& operator+=(const DataAndIndex& other) { + if (data > other.data) { + return *this; + } else { + data = other.data; + index = other.index; + return *this; + } + } +}; + +template +__device__ __forceinline__ vec_t GenerateGumbelNoise(uint64_t philox_seed, + uint64_t philox_offset, + uint64_t subsequence) { + curandStatePhilox4_32_10_t state; + vec_t noise; + constexpr float kEPSILON = 1e-20f; + constexpr float kLOG2 = 0.6931471806f; + auto uniform2gumbel = [](float x) { return -kLOG2 * log2f(-log2f(x + kEPSILON) + kEPSILON); }; +// TODO: compare the speed of log2 and log +#pragma unroll + for (uint32_t i = 0; i + 4 <= VEC_SIZE; i += 4) { + curand_init(philox_seed, subsequence + i, philox_offset, &state); + float4 noise_vec = curand_uniform4(&state); + noise[i] = uniform2gumbel(noise_vec.x); + noise[i + 1] = uniform2gumbel(noise_vec.y); + noise[i + 2] = uniform2gumbel(noise_vec.z); + noise[i + 3] = uniform2gumbel(noise_vec.w); + } + if constexpr (VEC_SIZE % 4 != 0) { + curand_init(philox_seed, subsequence + VEC_SIZE / 4 * 4, philox_offset, &state); + float4 noise_vec = curand_uniform4(&state); + if constexpr (VEC_SIZE % 4 == 1) { + noise[VEC_SIZE - 1] = uniform2gumbel(noise_vec.x); + } else if constexpr (VEC_SIZE % 4 == 2) { + noise[VEC_SIZE - 2] = uniform2gumbel(noise_vec.x); + noise[VEC_SIZE - 1] = uniform2gumbel(noise_vec.y); + } else if constexpr (VEC_SIZE % 4 == 3) { + noise[VEC_SIZE - 3] = uniform2gumbel(noise_vec.x); + noise[VEC_SIZE - 2] = uniform2gumbel(noise_vec.y); + noise[VEC_SIZE - 1] = uniform2gumbel(noise_vec.z); + } + } + + if constexpr (std::is_same_v) { + return noise; + } else { + vec_t ret; +#pragma unroll + for (uint32_t i = 0; i < VEC_SIZE; ++i) { + ret[i] = static_cast(noise[i]); + } + return ret; + } +} + +template +__global__ void SamplingFromLogitsKernel(DType* logits, IdType* output, IdType* indices, uint32_t d, + uint64_t philox_seed, uint64_t philox_offset) { + const uint32_t bx = blockIdx.x, tx = threadIdx.x; + const uint32_t row_idx = indices == nullptr ? bx : indices[bx]; + using SharedMem = typename BlockReduce, BLOCK_THREADS, + REDUCE_ALGORITHM>::TempStorage; + extern __shared__ __align__(alignof(SharedMem)) uint8_t smem_sampling_logit[]; + auto& temp_storage = reinterpret_cast(smem_sampling_logit); + + vec_t logits_vec; + DataAndIndex max_data = {-cuda::std::numeric_limits::infinity(), 0}; + for (uint32_t i = 0; i < ceil_div(d, BLOCK_THREADS * VEC_SIZE); ++i) { + logits_vec.fill(-cuda::std::numeric_limits::infinity()); + if ((i * BLOCK_THREADS + tx) * VEC_SIZE < d) { + logits_vec.cast_load(logits + row_idx * d + i * BLOCK_THREADS * VEC_SIZE + tx * VEC_SIZE); + } + + vec_t gumbel_noise = GenerateGumbelNoise( + philox_seed, philox_offset, + static_cast(bx * d + (i * BLOCK_THREADS + tx) * VEC_SIZE)); + DataAndIndex cur_data[VEC_SIZE]; +#pragma unroll + for (uint32_t j = 0; j < VEC_SIZE; ++j) { + cur_data[j].data = (i * BLOCK_THREADS + tx) * VEC_SIZE + j < d + ? logits_vec[j] + gumbel_noise[j] + : -cuda::std::numeric_limits::infinity(); + cur_data[j].index = (i * BLOCK_THREADS + tx) * VEC_SIZE + j; + } + + max_data += + BlockReduce, BLOCK_THREADS, REDUCE_ALGORITHM>(temp_storage) + .Sum(cur_data); + } + if (tx == 0) { + output[bx] = max_data.index; + } +} + +template +__global__ void SamplingFromProbKernel(DType* probs, IdType* output, IdType* indices, uint32_t d, + uint64_t philox_seed, uint64_t philox_offset) { + curandStatePhilox4_32_10_t state; + const uint32_t bx = blockIdx.x, tx = threadIdx.x; + curand_init(philox_seed, bx, philox_offset, &state); + const uint32_t row_idx = indices == nullptr ? bx : indices[bx]; + + extern __shared__ __align__( + alignof(SamplingTempStorage)) + uint8_t smem_sampling[]; + auto& temp_storage = + reinterpret_cast&>( + smem_sampling); + temp_storage.sampled_id = d; + __syncthreads(); + + vec_t probs_vec; + float aggregate(0); + float u = curand_uniform(&state); + +#pragma unroll 2 + for (uint32_t i = 0; i < ceil_div(d, BLOCK_THREADS * VEC_SIZE); ++i) { + probs_vec.fill(0); + if ((i * BLOCK_THREADS + tx) * VEC_SIZE < d) { + probs_vec.cast_load(probs + row_idx * d + i * BLOCK_THREADS * VEC_SIZE + tx * VEC_SIZE); + } + + DeviceSamplingFromProb( + i, d, [](float x) { return x > 0; }, u, probs_vec, aggregate, &temp_storage); + if (float(aggregate) > u) { + break; + } + } + int sampled_id = temp_storage.sampled_id; + if (sampled_id == d) { + // NOTE(Zihao): this would happen when u is very close to 1 + // and the sum of probabilities is smaller than u + // In this case, we use the last valid index as the sampled id + sampled_id = temp_storage.last_valid_id; + } + output[bx] = sampled_id; +} + +template +__global__ void TopKSamplingFromProbKernel(DType* probs, IdType* output, IdType* indices, + IdType* top_k_arr, uint32_t top_k_val, uint32_t d, + uint64_t philox_seed, uint64_t philox_offset) { + const uint32_t batch_size = gridDim.x; + const uint32_t bx = blockIdx.x, tx = threadIdx.x; + curandStatePhilox4_32_10_t state; + curand_init(philox_seed, bx, philox_offset, &state); + const uint32_t k = top_k_arr == nullptr ? top_k_val : top_k_arr[bx]; + const uint32_t row_idx = indices == nullptr ? bx : indices[bx]; + + extern __shared__ __align__( + alignof(SamplingTempStorage)) + uint8_t smem_sampling[]; + auto& temp_storage = + reinterpret_cast&>( + smem_sampling); + + vec_t probs_vec; + float aggregate; + float q = 1; + double low = 0, high = 1.f; + int sampled_id; + int round = 0; + do { + round += 1; + temp_storage.sampled_id = d; + __syncthreads(); + float u = curand_uniform(&state) * q; + aggregate = 0; +#pragma unroll 2 + for (uint32_t i = 0; i < ceil_div(d, BLOCK_THREADS * VEC_SIZE); ++i) { + probs_vec.fill(0); + if ((i * BLOCK_THREADS + tx) * VEC_SIZE < d) { + probs_vec.cast_load(probs + row_idx * d + (i * BLOCK_THREADS + tx) * VEC_SIZE); + } + + DeviceSamplingFromProb( + i, d, [&](float x) { return x > low; }, u, probs_vec, aggregate, &temp_storage); + if (aggregate > u) { + break; + } + } + __syncthreads(); + sampled_id = temp_storage.sampled_id; + if (sampled_id == d) { + // NOTE(Zihao): this would happen when u is very close to 1 + // and the sum of probabilities is smaller than u + // In this case, we use the last valid index as the sampled id + sampled_id = temp_storage.last_valid_id; + } + double pivot_0 = probs[row_idx * d + sampled_id]; + double pivot_1 = (pivot_0 + high) / 2; + + ValueCount aggregate_gt_pivot_0{0, 0}, aggregate_gt_pivot_1{0, 0}; +#pragma unroll 2 + for (uint32_t i = 0; i < ceil_div(d, BLOCK_THREADS * VEC_SIZE); ++i) { + probs_vec.fill(0); + if ((i * BLOCK_THREADS + tx) * VEC_SIZE < d) { + probs_vec.cast_load(probs + row_idx * d + (i * BLOCK_THREADS + tx) * VEC_SIZE); + } + + ValueCount probs_gt_pivot_0[VEC_SIZE], probs_gt_pivot_1[VEC_SIZE]; +#pragma unroll + for (uint32_t j = 0; j < VEC_SIZE; ++j) { + probs_gt_pivot_0[j] = { + (probs_vec[j] > pivot_0) ? probs_vec[j] : 0, + (probs_vec[j] > pivot_0 && (i * BLOCK_THREADS + tx) * VEC_SIZE + j < d)}; + probs_gt_pivot_1[j] = { + (probs_vec[j] > pivot_1) ? probs_vec[j] : 0, + (probs_vec[j] > pivot_1 && (i * BLOCK_THREADS + tx) * VEC_SIZE + j < d)}; + } + + aggregate_gt_pivot_0 += BlockReduce, BLOCK_THREADS, REDUCE_ALGORITHM>( + temp_storage.block_prim.reduce_value_count) + .Sum(probs_gt_pivot_0); + if (tx == 0) { + temp_storage.block_aggregate.pair = aggregate_gt_pivot_0; + } + __syncthreads(); + aggregate_gt_pivot_0 = temp_storage.block_aggregate.pair; + + aggregate_gt_pivot_1 += BlockReduce, BLOCK_THREADS, REDUCE_ALGORITHM>( + temp_storage.block_prim.reduce_value_count) + .Sum(probs_gt_pivot_1); + if (tx == 0) { + temp_storage.block_aggregate.pair = aggregate_gt_pivot_1; + } + __syncthreads(); + aggregate_gt_pivot_1 = temp_storage.block_aggregate.pair; + } + if (aggregate_gt_pivot_0.count < k) { + // case 1: pivot_0 accepted + break; + } + if (aggregate_gt_pivot_1.count < k) { + // case 2: pivot_0 rejected, pivot_1 accepted + low = pivot_0; + high = pivot_1; + q = aggregate_gt_pivot_0.value; + } else { + // case 3: pivot_0 rejected, pivot_1 rejected + low = pivot_1; + q = aggregate_gt_pivot_1.value; + } + } while (low < high); + __syncthreads(); + if (tx == 0) { + output[bx] = sampled_id; + } +} + +template +__global__ void TopPSamplingFromProbKernel(DType* probs, IdType* output, IdType* indices, + float* top_p_arr, float top_p_val, uint32_t d, + uint64_t philox_seed, uint64_t philox_offset) { + const uint32_t batch_size = gridDim.x; + const uint32_t bx = blockIdx.x, tx = threadIdx.x; + curandStatePhilox4_32_10_t state; + curand_init(philox_seed, bx, philox_offset, &state); + const uint32_t row_idx = indices == nullptr ? bx : indices[bx]; + float top_p = (top_p_arr == nullptr) ? top_p_val : top_p_arr[row_idx]; + + extern __shared__ __align__( + alignof(SamplingTempStorage)) + uint8_t smem_sampling[]; + auto& temp_storage = + reinterpret_cast&>( + smem_sampling); + + vec_t probs_vec; + float aggregate; + float q = 1; + double low = 0, high = 1.f; + int sampled_id; + do { + temp_storage.sampled_id = d; + __syncthreads(); + float u = curand_uniform(&state) * q; + aggregate = 0; +#pragma unroll 2 + for (uint32_t i = 0; i < ceil_div(d, BLOCK_THREADS * VEC_SIZE); ++i) { + probs_vec.fill(0); + if ((i * BLOCK_THREADS + tx) * VEC_SIZE < d) { + probs_vec.cast_load(probs + row_idx * d + (i * BLOCK_THREADS + tx) * VEC_SIZE); + } + + DeviceSamplingFromProb( + i, d, [&](float x) { return x > low; }, u, probs_vec, aggregate, &temp_storage); + if (aggregate > u) { + break; + } + } + __syncthreads(); + sampled_id = temp_storage.sampled_id; + if (sampled_id == d) { + // NOTE(Zihao): this would happen when u is very close to 1 + // and the sum of probabilities is smaller than u + // In this case, we use the last valid index as the sampled id + sampled_id = temp_storage.last_valid_id; + } + double pivot_0 = probs[row_idx * d + sampled_id]; + double pivot_1 = (pivot_0 + high) / 2; + + float aggregate_gt_pivot_0 = 0, aggregate_gt_pivot_1 = 0; +#pragma unroll 2 + for (uint32_t i = 0; i < ceil_div(d, BLOCK_THREADS * VEC_SIZE); ++i) { + probs_vec.fill(0); + if ((i * BLOCK_THREADS + tx) * VEC_SIZE < d) { + probs_vec.cast_load(probs + row_idx * d + (i * BLOCK_THREADS + tx) * VEC_SIZE); + } + + float probs_gt_pivot_0[VEC_SIZE], probs_gt_pivot_1[VEC_SIZE]; +#pragma unroll + for (uint32_t j = 0; j < VEC_SIZE; ++j) { + probs_gt_pivot_0[j] = (probs_vec[j] > pivot_0) ? probs_vec[j] : 0; + probs_gt_pivot_1[j] = (probs_vec[j] > pivot_1) ? probs_vec[j] : 0; + } + + aggregate_gt_pivot_0 += BlockReduce(temp_storage.block_prim.reduce) + .Sum(probs_gt_pivot_0); + if (tx == 0) { + temp_storage.block_aggregate.value = aggregate_gt_pivot_0; + } + __syncthreads(); + aggregate_gt_pivot_0 = temp_storage.block_aggregate.value; + + aggregate_gt_pivot_1 += BlockReduce(temp_storage.block_prim.reduce) + .Sum(probs_gt_pivot_1); + if (tx == 0) { + temp_storage.block_aggregate.value = aggregate_gt_pivot_1; + } + __syncthreads(); + aggregate_gt_pivot_1 = temp_storage.block_aggregate.value; + } + if (aggregate_gt_pivot_0 < top_p) { + // case 1: pivot_0 accepted + break; + } + if (aggregate_gt_pivot_1 < top_p) { + // case 2: pivot_0 rejected, pivot_1 accepted + low = pivot_0; + high = pivot_1; + q = aggregate_gt_pivot_0; + } else { + // case 3: pivot_0 rejected, pivot_1 rejected + low = pivot_1; + q = aggregate_gt_pivot_1; + } + } while (low < high); + __syncthreads(); + if (tx == 0) { + output[bx] = sampled_id; + } +} + +template +__global__ void MinPSamplingFromProbKernel(DType* probs, float* min_p_arr, IdType* output, + IdType* indices, float min_p_val, uint32_t d, + uint64_t philox_seed, uint64_t philox_offset) { + const uint32_t bx = blockIdx.x, tx = threadIdx.x; + float p = (min_p_arr == nullptr) ? min_p_val : min_p_arr[bx]; + curandStatePhilox4_32_10_t state; + curand_init(philox_seed, bx, philox_offset, &state); + const uint32_t row_idx = indices == nullptr ? bx : indices[bx]; + + extern __shared__ __align__( + alignof(SamplingTempStorage)) + uint8_t smem_sampling[]; + auto& temp_storage = + reinterpret_cast&>( + smem_sampling); + + float max_val = GetMaxValue>( + probs, row_idx, d, temp_storage); + float pivot = max_val * p; + + vec_t probs_vec; + float aggregate_gt_pivot = 0; +#pragma unroll 2 + for (uint32_t i = 0; i < ceil_div(d, BLOCK_THREADS * VEC_SIZE); ++i) { + probs_vec.fill(0); + if ((i * BLOCK_THREADS + tx) * VEC_SIZE < d) { + probs_vec.cast_load(probs + row_idx * d + (i * BLOCK_THREADS + tx) * VEC_SIZE); + } + + float probs_gt_pivot[VEC_SIZE]; +#pragma unroll + for (uint32_t j = 0; j < VEC_SIZE; ++j) { + probs_gt_pivot[j] = (probs_vec[j] >= pivot) ? probs_vec[j] : 0; + } + + aggregate_gt_pivot += BlockReduce(temp_storage.block_prim.reduce) + .Sum(probs_gt_pivot); + if (tx == 0) { + temp_storage.block_aggregate.value = aggregate_gt_pivot; + } + __syncthreads(); + } + + float aggregate = 0; + float q = temp_storage.block_aggregate.value; + + int sampled_id; + temp_storage.sampled_id = d; + __syncthreads(); + float u = curand_uniform(&state) * q; +#pragma unroll 2 + for (uint32_t i = 0; i < ceil_div(d, BLOCK_THREADS * VEC_SIZE); ++i) { + probs_vec.fill(0); + if ((i * BLOCK_THREADS + tx) * VEC_SIZE < d) { + probs_vec.cast_load(probs + row_idx * d + (i * BLOCK_THREADS + tx) * VEC_SIZE); + } + + DeviceSamplingFromProb( + i, d, [&](float x) { return x >= pivot; }, u, probs_vec, aggregate, &temp_storage); + if (aggregate > u) { + break; + } + } + sampled_id = temp_storage.sampled_id; + if (sampled_id == d) { + // NOTE(Zihao): this would happen when u is very close to 1 + // and the sum of probabilities is smaller than u + // In this case, we use the last valid index as the sampled id + sampled_id = temp_storage.last_valid_id; + } + output[bx] = sampled_id; +} + +template +__global__ void TopKTopPSamplingFromProbKernel(DType* probs, IdType* top_k_arr, float* top_p_arr, + IdType* output, IdType* indices, IdType top_k_val, + float top_p_val, uint32_t d, uint64_t philox_seed, + uint64_t philox_offset) { + const uint32_t batch_size = gridDim.x; + const uint32_t bx = blockIdx.x, tx = threadIdx.x; + curandStatePhilox4_32_10_t state; + curand_init(philox_seed, bx, philox_offset, &state); + const uint32_t row_idx = indices == nullptr ? bx : indices[bx]; + const uint32_t k = top_k_arr == nullptr ? top_k_val : top_k_arr[row_idx]; + const float p = top_p_arr == nullptr ? top_p_val : top_p_arr[row_idx]; + + extern __shared__ __align__( + alignof(SamplingTempStorage)) + uint8_t smem_sampling[]; + auto& temp_storage = + reinterpret_cast&>( + smem_sampling); + + vec_t probs_vec; + float aggregate; + float q = 1; + double low = 0, high = 1.f; + int sampled_id; + do { + temp_storage.sampled_id = d; + __syncthreads(); + float u = curand_uniform(&state) * q; + aggregate = 0; +#pragma unroll 2 + for (uint32_t i = 0; i < ceil_div(d, BLOCK_THREADS * VEC_SIZE); ++i) { + probs_vec.fill(0); + if ((i * BLOCK_THREADS + tx) * VEC_SIZE < d) { + probs_vec.cast_load(probs + row_idx * d + (i * BLOCK_THREADS + tx) * VEC_SIZE); + } + + DeviceSamplingFromProb( + i, d, [&](float x) { return x > low; }, u, probs_vec, aggregate, &temp_storage); + if (aggregate > u) { + break; + } + } + __syncthreads(); + sampled_id = temp_storage.sampled_id; + if (sampled_id == d) { + // NOTE(Zihao): this would happen when u is very close to 1 + // and the sum of probabilities is smaller than u + // In this case, we use the last valid index as the sampled id + sampled_id = temp_storage.last_valid_id; + } + double pivot_0 = probs[row_idx * d + sampled_id]; + double pivot_1 = (pivot_0 + high) / 2; + + ValueCount aggregate_gt_pivot_0{0, 0}, aggregate_gt_pivot_1{0, 0}; +#pragma unroll 2 + for (uint32_t i = 0; i < ceil_div(d, BLOCK_THREADS * VEC_SIZE); ++i) { + probs_vec.fill(0); + if ((i * BLOCK_THREADS + tx) * VEC_SIZE < d) { + probs_vec.cast_load(probs + row_idx * d + (i * BLOCK_THREADS + tx) * VEC_SIZE); + } + + ValueCount probs_gt_pivot_0[VEC_SIZE], probs_gt_pivot_1[VEC_SIZE]; +#pragma unroll + for (uint32_t j = 0; j < VEC_SIZE; ++j) { + probs_gt_pivot_0[j] = { + (probs_vec[j] > pivot_0) ? probs_vec[j] : 0, + (probs_vec[j] > pivot_0 && (i * BLOCK_THREADS + tx) * VEC_SIZE + j < d)}; + probs_gt_pivot_1[j] = { + (probs_vec[j] > pivot_1) ? probs_vec[j] : 0, + (probs_vec[j] > pivot_1 && (i * BLOCK_THREADS + tx) * VEC_SIZE + j < d)}; + } + + aggregate_gt_pivot_0 += + BlockReduce, BLOCK_THREADS>(temp_storage.block_prim.reduce_value_count) + .Sum(probs_gt_pivot_0); + if (tx == 0) { + temp_storage.block_aggregate.pair = aggregate_gt_pivot_0; + } + __syncthreads(); + aggregate_gt_pivot_0 = temp_storage.block_aggregate.pair; + + aggregate_gt_pivot_1 += + BlockReduce, BLOCK_THREADS>(temp_storage.block_prim.reduce_value_count) + .Sum(probs_gt_pivot_1); + if (tx == 0) { + temp_storage.block_aggregate.pair = aggregate_gt_pivot_1; + } + __syncthreads(); + aggregate_gt_pivot_1 = temp_storage.block_aggregate.pair; + } + if (aggregate_gt_pivot_0.count < k && aggregate_gt_pivot_0.value < p) { + // case 1: pivot_0 accepted + break; + } + if (aggregate_gt_pivot_1.count < k && aggregate_gt_pivot_1.value < p) { + // case 2: pivot_0 rejected, pivot_1 accepted + low = pivot_0; + high = pivot_1; + q = aggregate_gt_pivot_0.value; + } else { + // case 3: pivot_0 rejected, pivot_1 rejected + low = pivot_1; + q = aggregate_gt_pivot_1.value; + } + } while (low < high); + __syncthreads(); + if (tx == 0) { + output[bx] = sampled_id; + } +} + +template +cudaError_t OnlineSoftmax(DType* logits, DType* output, uint32_t batch_size, uint32_t d, + DType* temperature_arr, DType temperature_val, void* workspace_buffer, + size_t workspace_buffer_size_in_bytes, bool enable_pdl, + cudaStream_t stream = 0) { + constexpr uint32_t SMALL_BATCH_THRESHOLD = 128; + constexpr uint32_t LARGE_VOCAB_THRESHOLD = 24576; + constexpr uint32_t DEFAULT_SLICE_SIZE = 8192; + + const uint32_t vec_size = std::gcd(16 / sizeof(DType), d); + auto compute_capacity = GetCudaComputeCapability(); + + DISPATCH_COMPUTE_CAP_NUM_THREADS( + compute_capacity, BLOCK_THREADS, {DISPATCH_ALIGNED_VEC_SIZE(vec_size, VEC_SIZE, { + if (batch_size <= SMALL_BATCH_THRESHOLD && d >= LARGE_VOCAB_THRESHOLD) { + // Path A: Vocab-Splitting Strategy for small-batch & large-vocab + uint32_t num_slices = ceil_div(d, DEFAULT_SLICE_SIZE); + + const size_t partial_buffer_size = batch_size * num_slices * sizeof(PartialSoftmaxResult); + if (workspace_buffer_size_in_bytes < partial_buffer_size) { + return cudaErrorInvalidValue; + } + + AlignedAllocator allocator(workspace_buffer, workspace_buffer_size_in_bytes); + auto partial_results = allocator.aligned_alloc( + partial_buffer_size, alignof(PartialSoftmaxResult), "softmax_workspace"); + + // Phase 1: Map-Reduce across vocab slices + dim3 phase1_nblks(batch_size, num_slices); + dim3 phase1_nthrs(BLOCK_THREADS); + size_t smem_size = sizeof(OnlineSoftmaxTempStorage); + + auto phase1_kernel = OnlineSoftmaxMapKernel; + void* phase1_args[] = {&logits, &partial_results, &temperature_arr, &temperature_val, + &d, &num_slices}; + + FLASHINFER_CUDA_CALL(cudaFuncSetAttribute( + phase1_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + + if (enable_pdl) { + cudaLaunchAttribute attribute[1]; + attribute[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attribute[0].val.programmaticStreamSerializationAllowed = 1; + + cudaLaunchConfig_t config; + config.gridDim = phase1_nblks; + config.blockDim = phase1_nthrs; + config.dynamicSmemBytes = smem_size; + config.stream = stream; + config.attrs = attribute; + config.numAttrs = 1; + + FLASHINFER_CUDA_CALL(cudaLaunchKernelEx(&config, phase1_kernel, logits, partial_results, + temperature_arr, temperature_val, d, + num_slices)); + } else { + FLASHINFER_CUDA_CALL(cudaLaunchKernel((void*)phase1_kernel, phase1_nblks, phase1_nthrs, + phase1_args, smem_size, stream)); + } + + // Phase 2: Final reduction and apply normalization + dim3 phase2_nblks(batch_size); + dim3 phase2_nthrs(BLOCK_THREADS); + + auto phase2_kernel = OnlineSoftmaxReduceKernel; + void* phase2_args[] = {&logits, &output, &partial_results, &temperature_arr, + &temperature_val, &d, &num_slices}; + + FLASHINFER_CUDA_CALL(cudaFuncSetAttribute( + phase2_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + + if (enable_pdl) { + cudaLaunchAttribute attribute[1]; + attribute[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attribute[0].val.programmaticStreamSerializationAllowed = 1; + + cudaLaunchConfig_t config; + config.gridDim = phase2_nblks; + config.blockDim = phase2_nthrs; + config.dynamicSmemBytes = smem_size; + config.stream = stream; + config.attrs = attribute; + config.numAttrs = 1; + + FLASHINFER_CUDA_CALL(cudaLaunchKernelEx(&config, phase2_kernel, logits, output, + partial_results, temperature_arr, + temperature_val, d, num_slices)); + } else { + FLASHINFER_CUDA_CALL(cudaLaunchKernel((void*)phase2_kernel, phase2_nblks, phase2_nthrs, + phase2_args, smem_size, stream)); + } + } else { + // Path B: Single-Block Strategy + // Switch input cache + uint32_t cache_threshold; + if (batch_size <= 16) { + cache_threshold = 4096; + } else if (batch_size <= 32) { + cache_threshold = 2048; + } else { + cache_threshold = 0; + } + const bool cache_input = d <= cache_threshold; + + dim3 nblks(batch_size); + dim3 nthrs(BLOCK_THREADS); + void* args[] = {&logits, &output, &temperature_arr, &temperature_val, &d}; + + const size_t smem_logits_bytes = (round_up(d, VEC_SIZE) + VEC_SIZE) * sizeof(DType); + + uint32_t smem_size = sizeof(OnlineSoftmaxTempStorage) + + (cache_input ? smem_logits_bytes : 0); + + DISPATCH_SOFTMAX_CACHE_INPUT(cache_input, CACHE_INPUT, { + auto kernel = OnlineSoftmaxFusedKernel; + FLASHINFER_CUDA_CALL(cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + + if (enable_pdl) { + cudaLaunchAttribute attribute[1]; + attribute[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attribute[0].val.programmaticStreamSerializationAllowed = 1; + + cudaLaunchConfig_t config; + config.gridDim = nblks; + config.blockDim = nthrs; + config.dynamicSmemBytes = smem_size; + config.stream = stream; + config.attrs = attribute; + config.numAttrs = 1; + + FLASHINFER_CUDA_CALL(cudaLaunchKernelEx(&config, kernel, logits, output, + temperature_arr, temperature_val, d)); + } else { + FLASHINFER_CUDA_CALL( + cudaLaunchKernel((void*)kernel, nblks, nthrs, args, smem_size, stream)); + } + }); + } + })}); + return cudaSuccess; +} + +template +cudaError_t SamplingFromLogits(T* logits, IdType* output, IdType* indices, uint32_t batch_size, + uint32_t d, bool deterministic, uint64_t philox_seed, + uint64_t philox_offset, cudaStream_t stream = 0) { + constexpr uint32_t BLOCK_THREADS = 1024; + const uint32_t vec_size = std::gcd(16 / sizeof(T), d); + dim3 nblks(batch_size); + dim3 nthrs(BLOCK_THREADS); + void* args[] = {&logits, &output, &indices, &d, &philox_seed, &philox_offset}; + const uint32_t smem_size = sizeof( + typename BlockReduce, BLOCK_THREADS, REDUCE_ALGO>::TempStorage); + + DISPATCH_ALIGNED_VEC_SIZE( + vec_size, VEC_SIZE, {DISPATCH_DETERMINISTIC(deterministic, DETERMINISTIC, { + auto kernel = SamplingFromLogitsKernel; + FLASHINFER_CUDA_CALL( + cudaLaunchKernel((void*)kernel, nblks, nthrs, args, smem_size, stream)); + })}); + return cudaSuccess; +} + +template +cudaError_t SamplingFromProb(T* probs, IdType* output, IdType* indices, uint32_t batch_size, + uint32_t d, bool deterministic, uint64_t philox_seed, + uint64_t philox_offset, cudaStream_t stream = 0) { + constexpr uint32_t BLOCK_THREADS = 1024; + const uint32_t vec_size = std::gcd(16 / sizeof(T), d); + dim3 nblks(batch_size); + dim3 nthrs(BLOCK_THREADS); + void* args[] = {&probs, &output, &indices, &d, &philox_seed, &philox_offset, &d}; + const uint32_t smem_size = sizeof(SamplingTempStorage); + + DISPATCH_ALIGNED_VEC_SIZE( + vec_size, VEC_SIZE, {DISPATCH_DETERMINISTIC(deterministic, DETERMINISTIC, { + auto kernel = SamplingFromProbKernel; + FLASHINFER_CUDA_CALL( + cudaLaunchKernel((void*)kernel, nblks, nthrs, args, smem_size, stream)); + })}); + return cudaSuccess; +} + +template +cudaError_t TopKSamplingFromProb(T* probs, IdType* output, IdType* indices, T* top_k_arr, + uint32_t batch_size, uint32_t top_k_val, uint32_t d, + bool deterministic, uint64_t philox_seed, uint64_t philox_offset, + cudaStream_t stream = 0) { + const uint32_t vec_size = std::gcd(16 / sizeof(T), d); + + auto compute_capacity = GetCudaComputeCapability(); + DISPATCH_COMPUTE_CAP_NUM_THREADS(compute_capacity, BLOCK_THREADS, { + const uint32_t smem_size = sizeof(SamplingTempStorage); + dim3 nblks(batch_size); + dim3 nthrs(BLOCK_THREADS); + void* args[] = {&probs, &output, &indices, &top_k_arr, + &top_k_val, &d, &philox_seed, &philox_offset}; + + DISPATCH_ALIGNED_VEC_SIZE( + vec_size, VEC_SIZE, {DISPATCH_DETERMINISTIC(deterministic, DETERMINISTIC, { + auto kernel = TopKSamplingFromProbKernel; + FLASHINFER_CUDA_CALL( + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + FLASHINFER_CUDA_CALL( + cudaLaunchKernel((void*)kernel, nblks, nthrs, args, smem_size, stream)); + })}); + return cudaSuccess; + }); +} + +template +cudaError_t TopPSamplingFromProb(T* probs, IdType* output, IdType* indices, T* top_p_arr, + uint32_t batch_size, T top_p_val, uint32_t d, bool deterministic, + uint64_t philox_seed, uint64_t philox_offset, + cudaStream_t stream = 0) { + constexpr uint32_t BLOCK_THREADS = 1024; + const uint32_t vec_size = std::gcd(16 / sizeof(T), d); + + const uint32_t smem_size = sizeof(SamplingTempStorage); + dim3 nblks(batch_size); + dim3 nthrs(BLOCK_THREADS); + void* args[] = {&probs, &output, &indices, &top_p_arr, + &top_p_val, &d, &philox_seed, &philox_offset}; + + DISPATCH_ALIGNED_VEC_SIZE( + vec_size, VEC_SIZE, {DISPATCH_DETERMINISTIC(deterministic, DETERMINISTIC, { + auto kernel = TopPSamplingFromProbKernel; + FLASHINFER_CUDA_CALL( + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + FLASHINFER_CUDA_CALL( + cudaLaunchKernel((void*)kernel, nblks, nthrs, args, smem_size, stream)); + })}); + return cudaSuccess; +} + +template +cudaError_t MinPSamplingFromProb(T* probs, T* min_p_arr, IdType* output, IdType* indices, + uint32_t batch_size, float min_p_val, uint32_t d, + bool deterministic, uint64_t philox_seed, uint64_t philox_offset, + cudaStream_t stream = 0) { + constexpr uint32_t BLOCK_THREADS = 1024; + const uint32_t vec_size = std::gcd(16 / sizeof(T), d); + + const uint32_t smem_size = sizeof(SamplingTempStorage); + dim3 nblks(batch_size); + dim3 nthrs(BLOCK_THREADS); + void* args[] = {&probs, &min_p_arr, &output, &indices, + &min_p_val, &d, &philox_seed, &philox_offset}; + + DISPATCH_ALIGNED_VEC_SIZE( + vec_size, VEC_SIZE, {DISPATCH_DETERMINISTIC(deterministic, DETERMINISTIC, { + auto kernel = MinPSamplingFromProbKernel; + FLASHINFER_CUDA_CALL( + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + FLASHINFER_CUDA_CALL( + cudaLaunchKernel((void*)kernel, nblks, nthrs, args, smem_size, stream)); + })}); + return cudaSuccess; +} + +template +cudaError_t TopKTopPSamplingFromProb(T* probs, IdType* top_k_arr, T* top_p_arr, IdType* output, + IdType* indices, uint32_t batch_size, IdType top_k_val, + T top_p_val, uint32_t d, bool deterministic, + uint64_t philox_seed, uint64_t philox_offset, + cudaStream_t stream = 0) { + const uint32_t vec_size = std::gcd(16 / sizeof(T), d); + + auto compute_capacity = GetCudaComputeCapability(); + DISPATCH_COMPUTE_CAP_NUM_THREADS(compute_capacity, BLOCK_THREADS, { + const uint32_t smem_size = sizeof(SamplingTempStorage); + dim3 nblks(batch_size); + dim3 nthrs(BLOCK_THREADS); + void* args[] = {&probs, &top_k_arr, &top_p_arr, &output, &indices, + &top_k_val, &top_p_val, &d, &philox_seed, &philox_offset}; + + DISPATCH_ALIGNED_VEC_SIZE( + vec_size, VEC_SIZE, {DISPATCH_DETERMINISTIC(deterministic, DETERMINISTIC, { + auto kernel = TopKTopPSamplingFromProbKernel; + FLASHINFER_CUDA_CALL( + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + FLASHINFER_CUDA_CALL( + cudaLaunchKernel((void*)kernel, nblks, nthrs, args, smem_size, stream)); + })}); + return cudaSuccess; + }); +} + +template +struct RenormTempStorage { + union { + typename BlockReduce::TempStorage reduce; + typename BlockReduce::TempStorage reduce_int; + typename BlockReduce, BLOCK_THREADS, REDUCE_ALGORITHM>::TempStorage + reduce_value_count; + } block_prim; + struct { + float max_val; + float min_val; + union { + struct { + float values[2]; + }; + struct { + int counts[2]; + }; + struct { + ValueCount pairs[2]; + }; + } block_aggregate; + }; +}; + +template +__global__ void TopPRenormProbKernel(DType* probs, DType* renormed_prob, float* top_p_arr, + float top_p_val, uint32_t d) { + const uint32_t bx = blockIdx.x, tx = threadIdx.x; + const uint32_t row_idx = bx; + float p = top_p_arr == nullptr ? top_p_val : top_p_arr[bx]; + + extern __shared__ __align__(alignof(RenormTempStorage)) + uint8_t smem_renorm[]; + auto& temp_storage = + reinterpret_cast&>(smem_renorm); + temp_storage.max_val = 0; + vec_t probs_vec; + + float max_val = GetMaxValue>(probs, row_idx, d, + temp_storage); + + double low = 0, high = max_val; + float min_gt_low, max_le_high; + float sum_low = 1; + // f(x) = sum(probs[probs > x]), f(x) is non-increasing + // min_gt_low = min{p \in probs | p > low}, max_le_high = max{p \in probs | p <= high} + // loop invariant: + // - f(low) >= p, f(high) < p + // - f(low) > f(min_gt_low) >= f(max_le_high) == f(high) + // stopping condition + // - f(low) >= p, f(min_gt_low) == f(max_le_high) == f(high) < p + do { + double pivot_0 = (high + 2 * low) / 3; + double pivot_1 = (2 * high + low) / 3; + + float aggregate_gt_pivot_0 = 0, aggregate_gt_pivot_1 = 0; + min_gt_low = high; + max_le_high = low; +#pragma unroll 2 + for (uint32_t i = 0; i < ceil_div(d, BLOCK_THREADS * VEC_SIZE); ++i) { + probs_vec.fill(0); + if ((i * BLOCK_THREADS + tx) * VEC_SIZE < d) { + probs_vec.cast_load(probs + row_idx * d + i * BLOCK_THREADS * VEC_SIZE + tx * VEC_SIZE); + } + + float probs_gt_pivot_0[VEC_SIZE], probs_gt_pivot_1[VEC_SIZE]; +#pragma unroll + for (uint32_t j = 0; j < VEC_SIZE; ++j) { + probs_gt_pivot_0[j] = (probs_vec[j] > pivot_0) ? probs_vec[j] : 0; + probs_gt_pivot_1[j] = (probs_vec[j] > pivot_1) ? probs_vec[j] : 0; + + if (probs_vec[j] > low && (i * BLOCK_THREADS + tx) * VEC_SIZE + j < d) { + min_gt_low = min(min_gt_low, probs_vec[j]); + } + if (probs_vec[j] <= high && (i * BLOCK_THREADS + tx) * VEC_SIZE + j < d) { + max_le_high = max(max_le_high, probs_vec[j]); + } + } + + aggregate_gt_pivot_0 += + BlockReduce(temp_storage.block_prim.reduce) + .Sum(probs_gt_pivot_0); + __syncthreads(); + + aggregate_gt_pivot_1 += + BlockReduce(temp_storage.block_prim.reduce) + .Sum(probs_gt_pivot_1); + __syncthreads(); + } + min_gt_low = BlockReduce(temp_storage.block_prim.reduce) + .Reduce(min_gt_low, cub::Min()); + __syncthreads(); + max_le_high = + BlockReduce(temp_storage.block_prim.reduce) + .Reduce(max_le_high, cub::Max()); + if (tx == 0) { + temp_storage.block_aggregate.values[0] = aggregate_gt_pivot_0; + temp_storage.block_aggregate.values[1] = aggregate_gt_pivot_1; + temp_storage.min_val = min_gt_low; + temp_storage.max_val = max_le_high; + } + __syncthreads(); + aggregate_gt_pivot_0 = temp_storage.block_aggregate.values[0]; + aggregate_gt_pivot_1 = temp_storage.block_aggregate.values[1]; + min_gt_low = temp_storage.min_val; + max_le_high = temp_storage.max_val; + + if (aggregate_gt_pivot_1 >= p) { + low = pivot_1; + sum_low = aggregate_gt_pivot_1; + } else if (aggregate_gt_pivot_0 >= p) { + low = pivot_0; + high = min(pivot_1, max_le_high); + sum_low = aggregate_gt_pivot_0; + } else { + high = min(pivot_0, max_le_high); + } + } while (min_gt_low != max_le_high); + + float normalizer = math::ptx_rcp(max(sum_low, 1e-8)); + + // normalize +#pragma unroll 2 + for (uint32_t i = 0; i < ceil_div(d, BLOCK_THREADS * VEC_SIZE); ++i) { + probs_vec.fill(0); + if ((i * BLOCK_THREADS + tx) * VEC_SIZE < d) { + probs_vec.cast_load(probs + row_idx * d + i * BLOCK_THREADS * VEC_SIZE + tx * VEC_SIZE); + } +#pragma unroll + for (uint32_t j = 0; j < VEC_SIZE; ++j) { + probs_vec[j] = (probs_vec[j] > low) ? probs_vec[j] * normalizer : 0; + } + if ((i * BLOCK_THREADS + tx) * VEC_SIZE < d) { + probs_vec.cast_store(renormed_prob + row_idx * d + i * BLOCK_THREADS * VEC_SIZE + + tx * VEC_SIZE); + } + } +} + +template +__global__ void TopKMaskLogitsKernel(DType* logits, DType* masked_logits, IdType* top_k_arr, + uint32_t top_k_val, uint32_t d) { + const uint32_t bx = blockIdx.x, tx = threadIdx.x; + const uint32_t row_idx = bx; + uint32_t k = top_k_arr == nullptr ? top_k_val : top_k_arr[bx]; + double pivot = -cuda::std::numeric_limits::infinity(); + vec_t logits_vec; + if (k < d) { + extern __shared__ __align__(alignof(RenormTempStorage)) + uint8_t smem_renorm[]; + auto& temp_storage = + reinterpret_cast&>(smem_renorm); + float logits_greater_than_pivot[VEC_SIZE]; // pivot initialized to 0 + + auto [min_val, max_val] = GetMinMaxValue>( + logits, row_idx, d, temp_storage); + + double low = (min_val == -cuda::std::numeric_limits::infinity()) + ? cuda::std::numeric_limits::lowest() + : min_val - 1, + high = max_val; + float min_gt_low, max_le_high; + // f(x) = len(nonzero(probs > x)), f(x) is non-increasing + // min_gt_low = min{p \in probs | p > low}, max_le_high = max{p \in probs | p <= high} + // loop invariant: + // - f(low) >= k, f(high) < k + // - f(low) > f(min_gt_low) >= f(max_le_high) == f(high) + // stopping condition: min_gt_low == max_le_high + // - f(low) >= k, f(min_gt_low) == f(max_le_high) == f(high) < k + do { + double pivot_0 = (high + 2 * low) / 3; + double pivot_1 = (2 * high + low) / 3; + + int aggregate_gt_pivot_0 = 0, aggregate_gt_pivot_1 = 0; + min_gt_low = high; + max_le_high = low; +#pragma unroll 2 + for (uint32_t i = 0; i < ceil_div(d, BLOCK_THREADS * VEC_SIZE); ++i) { + logits_vec.fill(0); + if ((i * BLOCK_THREADS + tx) * VEC_SIZE < d) { + logits_vec.cast_load(logits + row_idx * d + i * BLOCK_THREADS * VEC_SIZE + tx * VEC_SIZE); + } + int probs_gt_pivot_0_count[VEC_SIZE], probs_gt_pivot_1_count[VEC_SIZE]; +#pragma unroll + for (uint32_t j = 0; j < VEC_SIZE; ++j) { + probs_gt_pivot_0_count[j] = + logits_vec[j] > pivot_0 && (i * BLOCK_THREADS + tx) * VEC_SIZE + j < d; + probs_gt_pivot_1_count[j] = + logits_vec[j] > pivot_1 && (i * BLOCK_THREADS + tx) * VEC_SIZE + j < d; + + if (logits_vec[j] > low && (i * BLOCK_THREADS + tx) * VEC_SIZE + j < d) { + min_gt_low = min(min_gt_low, logits_vec[j]); + } + if (logits_vec[j] <= high && (i * BLOCK_THREADS + tx) * VEC_SIZE + j < d) { + max_le_high = max(max_le_high, logits_vec[j]); + } + } + + aggregate_gt_pivot_0 += + BlockReduce(temp_storage.block_prim.reduce_int) + .Sum(probs_gt_pivot_0_count); + __syncthreads(); + + aggregate_gt_pivot_1 += + BlockReduce(temp_storage.block_prim.reduce_int) + .Sum(probs_gt_pivot_1_count); + __syncthreads(); + } + min_gt_low = + BlockReduce(temp_storage.block_prim.reduce) + .Reduce(min_gt_low, cub::Min()); + __syncthreads(); + max_le_high = + BlockReduce(temp_storage.block_prim.reduce) + .Reduce(max_le_high, cub::Max()); + if (tx == 0) { + temp_storage.block_aggregate.counts[0] = aggregate_gt_pivot_0; + temp_storage.block_aggregate.counts[1] = aggregate_gt_pivot_1; + temp_storage.min_val = min_gt_low; + temp_storage.max_val = max_le_high; + } + __syncthreads(); + aggregate_gt_pivot_0 = temp_storage.block_aggregate.counts[0]; + aggregate_gt_pivot_1 = temp_storage.block_aggregate.counts[1]; + min_gt_low = temp_storage.min_val; + max_le_high = temp_storage.max_val; + + if (aggregate_gt_pivot_1 >= k) { + low = pivot_1; + } else if (aggregate_gt_pivot_0 >= k) { + low = pivot_0; + high = min(pivot_1, max_le_high); + } else { + high = min(pivot_0, max_le_high); + } + } while (min_gt_low != max_le_high); + pivot = low; + } + + // masking +#pragma unroll 2 + for (uint32_t i = 0; i < ceil_div(d, BLOCK_THREADS * VEC_SIZE); ++i) { + logits_vec.fill(0); + if ((i * BLOCK_THREADS + tx) * VEC_SIZE < d) { + logits_vec.cast_load(logits + row_idx * d + i * BLOCK_THREADS * VEC_SIZE + tx * VEC_SIZE); + } +#pragma unroll + for (uint32_t j = 0; j < VEC_SIZE; ++j) { + logits_vec[j] = + (logits_vec[j] > pivot) ? logits_vec[j] : -cuda::std::numeric_limits::infinity(); + } + if ((i * BLOCK_THREADS + tx) * VEC_SIZE < d) { + logits_vec.store(masked_logits + row_idx * d + i * BLOCK_THREADS * VEC_SIZE + tx * VEC_SIZE); + } + } +} + +template +__global__ void TopKRenormProbKernel(DType* probs, DType* renormed_prob, IdType* top_k_arr, + uint32_t top_k_val, uint32_t d) { + const uint32_t bx = blockIdx.x, tx = threadIdx.x; + const uint32_t row_idx = bx; + uint32_t k = top_k_arr == nullptr ? top_k_val : top_k_arr[bx]; + double pivot = -cuda::std::numeric_limits::infinity(), normalizer = 1; + vec_t probs_vec; + if (k < d) { + extern __shared__ __align__(alignof(RenormTempStorage)) + uint8_t smem_renorm[]; + auto& temp_storage = + reinterpret_cast&>(smem_renorm); + temp_storage.max_val = 0; + + float max_val = GetMaxValue>( + probs, row_idx, d, temp_storage); + + double low = 0, high = max_val; + float min_gt_low, max_le_high; + float sum_low = 1; + // f(x) = len(nonzero(probs > x)), f(x) is non-increasing + // min_gt_low = min{p \in probs | p > low}, max_le_high = max{p \in probs | p <= high} + // loop invariant: + // - f(low) >= k, f(high) < k + // - f(low) > f(min_gt_low) >= f(max_le_high) == f(high) + // stopping condition: min_gt_low == max_le_high + // - f(low) >= k, f(min_gt_low) == f(max_le_high) == f(high) < k + do { + double pivot_0 = (high + 2 * low) / 3; + double pivot_1 = (2 * high + low) / 3; + + ValueCount aggregate_gt_pivot_0{0, 0}, aggregate_gt_pivot_1{0, 0}; + min_gt_low = high; + max_le_high = low; +#pragma unroll 2 + for (uint32_t i = 0; i < ceil_div(d, BLOCK_THREADS * VEC_SIZE); ++i) { + probs_vec.fill(0); + if ((i * BLOCK_THREADS + tx) * VEC_SIZE < d) { + probs_vec.cast_load(probs + row_idx * d + i * BLOCK_THREADS * VEC_SIZE + tx * VEC_SIZE); + } + ValueCount probs_gt_pivot_0_pair[VEC_SIZE], probs_gt_pivot_1_pair[VEC_SIZE]; +#pragma unroll + for (uint32_t j = 0; j < VEC_SIZE; ++j) { + probs_gt_pivot_0_pair[j] = { + (probs_vec[j] > pivot_0) ? probs_vec[j] : 0, + (probs_vec[j] > pivot_0 && (i * BLOCK_THREADS + tx) * VEC_SIZE + j < d)}; + probs_gt_pivot_1_pair[j] = { + (probs_vec[j] > pivot_1) ? probs_vec[j] : 0, + (probs_vec[j] > pivot_1 && (i * BLOCK_THREADS + tx) * VEC_SIZE + j < d)}; + + if (probs_vec[j] > low && (i * BLOCK_THREADS + tx) * VEC_SIZE + j < d) { + min_gt_low = min(min_gt_low, probs_vec[j]); + } + if (probs_vec[j] <= high && (i * BLOCK_THREADS + tx) * VEC_SIZE + j < d) { + max_le_high = max(max_le_high, probs_vec[j]); + } + } + + aggregate_gt_pivot_0 += BlockReduce, BLOCK_THREADS, REDUCE_ALGORITHM>( + temp_storage.block_prim.reduce_value_count) + .Sum(probs_gt_pivot_0_pair); + __syncthreads(); + + aggregate_gt_pivot_1 += BlockReduce, BLOCK_THREADS, REDUCE_ALGORITHM>( + temp_storage.block_prim.reduce_value_count) + .Sum(probs_gt_pivot_1_pair); + __syncthreads(); + } + min_gt_low = + BlockReduce(temp_storage.block_prim.reduce) + .Reduce(min_gt_low, cub::Min()); + __syncthreads(); + max_le_high = + BlockReduce(temp_storage.block_prim.reduce) + .Reduce(max_le_high, cub::Max()); + if (tx == 0) { + temp_storage.block_aggregate.pairs[0] = aggregate_gt_pivot_0; + temp_storage.block_aggregate.pairs[1] = aggregate_gt_pivot_1; + temp_storage.min_val = min_gt_low; + temp_storage.max_val = max_le_high; + } + __syncthreads(); + aggregate_gt_pivot_0 = temp_storage.block_aggregate.pairs[0]; + aggregate_gt_pivot_1 = temp_storage.block_aggregate.pairs[1]; + min_gt_low = temp_storage.min_val; + max_le_high = temp_storage.max_val; + + if (aggregate_gt_pivot_1.count >= k) { + low = pivot_1; + sum_low = float(aggregate_gt_pivot_1.value); + } else if (aggregate_gt_pivot_0.count >= k) { + low = pivot_0; + high = min(pivot_1, max_le_high); + sum_low = float(aggregate_gt_pivot_0.value); + } else { + high = min(pivot_0, max_le_high); + } + } while (min_gt_low != max_le_high); + + normalizer = math::ptx_rcp(max(sum_low, 1e-8)); + pivot = low; + } + + // normalize +#pragma unroll 2 + for (uint32_t i = 0; i < ceil_div(d, BLOCK_THREADS * VEC_SIZE); ++i) { + probs_vec.fill(0); + if ((i * BLOCK_THREADS + tx) * VEC_SIZE < d) { + probs_vec.cast_load(probs + row_idx * d + i * BLOCK_THREADS * VEC_SIZE + tx * VEC_SIZE); + } +#pragma unroll + for (uint32_t j = 0; j < VEC_SIZE; ++j) { + probs_vec[j] = (probs_vec[j] > pivot) ? probs_vec[j] * normalizer : 0; + } + if ((i * BLOCK_THREADS + tx) * VEC_SIZE < d) { + probs_vec.store(renormed_prob + row_idx * d + i * BLOCK_THREADS * VEC_SIZE + tx * VEC_SIZE); + } + } +} + +template +cudaError_t TopPRenormProb(DType* probs, DType* renormed_prob, float* top_p_arr, + uint32_t batch_size, float top_p_val, uint32_t d, + cudaStream_t stream = 0) { + constexpr uint32_t BLOCK_THREADS = 1024; + const uint32_t vec_size = std::gcd(16 / sizeof(DType), d); + + const uint32_t smem_size = sizeof(RenormTempStorage); + dim3 nblks(batch_size); + dim3 nthrs(BLOCK_THREADS); + void* args[] = {&probs, &renormed_prob, &top_p_arr, &top_p_val, &d}; + DISPATCH_ALIGNED_VEC_SIZE(vec_size, VEC_SIZE, { + auto kernel = TopPRenormProbKernel; + FLASHINFER_CUDA_CALL( + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + FLASHINFER_CUDA_CALL(cudaLaunchKernel((void*)kernel, nblks, nthrs, args, smem_size, stream)); + }); + return cudaSuccess; +} + +template +cudaError_t TopKRenormProb(DType* probs, DType* renormed_prob, IdType* top_k_arr, + uint32_t batch_size, uint32_t top_k_val, uint32_t d, + cudaStream_t stream = 0) { + const uint32_t vec_size = std::gcd(16 / sizeof(DType), d); + + auto compute_capacity = GetCudaComputeCapability(); + DISPATCH_COMPUTE_CAP_NUM_THREADS(compute_capacity, BLOCK_THREADS, { + const uint32_t smem_size = sizeof(RenormTempStorage); + dim3 nblks(batch_size); + dim3 nthrs(BLOCK_THREADS); + void* args[] = {&probs, &renormed_prob, &top_k_arr, &top_k_val, &d}; + DISPATCH_ALIGNED_VEC_SIZE(vec_size, VEC_SIZE, { + auto kernel = TopKRenormProbKernel; + FLASHINFER_CUDA_CALL( + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + FLASHINFER_CUDA_CALL(cudaLaunchKernel((void*)kernel, nblks, nthrs, args, smem_size, stream)); + }); + return cudaSuccess; + }); +} + +template +cudaError_t TopKMaskLogits(DType* logits, DType* masked_logits, IdType* top_k_arr, + uint32_t batch_size, uint32_t top_k_val, uint32_t d, + cudaStream_t stream = 0) { + const uint32_t vec_size = std::gcd(16 / sizeof(DType), d); + + auto compute_capacity = GetCudaComputeCapability(); + DISPATCH_COMPUTE_CAP_NUM_THREADS(compute_capacity, BLOCK_THREADS, { + const uint32_t smem_size = sizeof(RenormTempStorage); + dim3 nblks(batch_size); + dim3 nthrs(BLOCK_THREADS); + void* args[] = {&logits, &masked_logits, &top_k_arr, &top_k_val, &d}; + DISPATCH_ALIGNED_VEC_SIZE(vec_size, VEC_SIZE, { + auto kernel = TopKMaskLogitsKernel; + FLASHINFER_CUDA_CALL( + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + FLASHINFER_CUDA_CALL(cudaLaunchKernel((void*)kernel, nblks, nthrs, args, smem_size, stream)); + }); + return cudaSuccess; + }); +} + +template +__global__ void ChainSpeculativeSampling(DType* draft_probs, IdType* draft_token_ids, + DType* target_probs, IdType* output_token_ids, + IdType* output_accepted_token_num, + IdType* output_emitted_draft_token_num, + uint32_t num_speculative_tokens, uint32_t d, + uint64_t philox_seed, uint64_t philox_offset) { + const uint32_t bx = blockIdx.x, tx = threadIdx.x; + const uint32_t row_idx = bx; + curandStatePhilox4_32_10_t curand_state; + curand_init(philox_seed, bx, philox_offset, &curand_state); + + extern __shared__ __align__( + alignof(SamplingTempStorage)) + uint8_t smem_sampling[]; + auto& temp_storage = + reinterpret_cast&>( + smem_sampling); + + uint32_t pos = num_speculative_tokens; + for (uint32_t i = 0; i < num_speculative_tokens; ++i) { + IdType draft_id = draft_token_ids[row_idx * num_speculative_tokens + i]; + float q = target_probs[(row_idx * (num_speculative_tokens + 1) + i) * d + draft_id], + p = draft_probs[(row_idx * num_speculative_tokens + i) * d + draft_id]; + float u = curand_uniform(&curand_state); + if (u * p < q) { + // accept the draft models output + output_token_ids[row_idx * (num_speculative_tokens + 1) + i] = draft_id; + } else { + pos = i; + break; + } + } + + uint32_t emitted_token_num = pos; + uint32_t accepted_token_num = pos; + for (uint32_t i = pos; i < num_speculative_tokens; ++i) { + int draft_id = draft_token_ids[row_idx * num_speculative_tokens + i]; + float q = target_probs[(row_idx * (num_speculative_tokens + 1) + i) * d + draft_id], + p = draft_probs[(row_idx * num_speculative_tokens + i) * d + draft_id]; + float u = curand_uniform(&curand_state); + if (u * p < q) { + ++accepted_token_num; + } + } + + if (tx == 0) { + output_accepted_token_num[row_idx] += accepted_token_num; + output_emitted_draft_token_num[row_idx] += emitted_token_num; + } + + // sample from relu(target_probs - draft_probs) + float sum_relu_q_minus_p = 0; + vec_t q_vec, p_vec; + float relu_q_minus_p[VEC_SIZE]; +#pragma unroll 2 + for (uint32_t i = 0; i < ceil_div(d, BLOCK_THREADS * VEC_SIZE); ++i) { + q_vec.fill(0); + p_vec.fill(0); + if ((i * BLOCK_THREADS + tx) * VEC_SIZE < d) { + q_vec.cast_load(target_probs + (row_idx * (num_speculative_tokens + 1) + pos) * d + + i * BLOCK_THREADS * VEC_SIZE + tx * VEC_SIZE); + if (pos != num_speculative_tokens) { + // there is no draft_probs for the bonus token + p_vec.cast_load(draft_probs + (row_idx * num_speculative_tokens + pos) * d + + i * BLOCK_THREADS * VEC_SIZE + tx * VEC_SIZE); + } + } +#pragma unroll + for (uint32_t j = 0; j < VEC_SIZE; ++j) { + relu_q_minus_p[j] = max(q_vec[j] - p_vec[j], 0); + } + sum_relu_q_minus_p += + BlockReduce(temp_storage.block_prim.reduce) + .Sum(relu_q_minus_p); + __syncthreads(); + } + if (tx == 0) { + temp_storage.block_aggregate.value = sum_relu_q_minus_p; + } + // init the first rejected token to d + temp_storage.sampled_id = d; + __syncthreads(); + sum_relu_q_minus_p = temp_storage.block_aggregate.value; + float u = curand_uniform(&curand_state) * sum_relu_q_minus_p; + + float aggregate_relu_q_minus_p(0); +#pragma unroll 2 + for (uint32_t i = 0; i < ceil_div(d, BLOCK_THREADS * VEC_SIZE); ++i) { + q_vec.fill(0); + p_vec.fill(0); + if ((i * BLOCK_THREADS + tx) * VEC_SIZE < d) { + q_vec.cast_load(target_probs + (row_idx * (num_speculative_tokens + 1) + pos) * d + + i * BLOCK_THREADS * VEC_SIZE + tx * VEC_SIZE); + if (pos != num_speculative_tokens) { + // there is no draft_probs for the bonus token + p_vec.cast_load(draft_probs + (row_idx * num_speculative_tokens + pos) * d + + i * BLOCK_THREADS * VEC_SIZE + tx * VEC_SIZE); + } + } + + vec_t relu_q_minus_p_vec; +#pragma unroll + for (uint32_t j = 0; j < VEC_SIZE; ++j) { + relu_q_minus_p_vec[j] = max(q_vec[j] - p_vec[j], 0); + } + + DeviceSamplingFromProb( + i, d, [&](float x) { return x > 0; }, u, relu_q_minus_p_vec, aggregate_relu_q_minus_p, + &temp_storage); + if (aggregate_relu_q_minus_p > u) { + break; + } + } + __syncthreads(); + int sampled_id = temp_storage.sampled_id; + if (sampled_id == d) { + // NOTE(Zihao): this would happen when u is very close to 1 + // and the sum of probabilities is smaller than u + // In this case, we use the last valid index as the sampled id + sampled_id = temp_storage.last_valid_id; + } + // set the first rejected token + output_token_ids[row_idx * (num_speculative_tokens + 1) + pos] = sampled_id; + // move to the next token + pos++; + + // pad remaining tokens with -1 + for (; pos < num_speculative_tokens + 1; ++pos) { + output_token_ids[row_idx * (num_speculative_tokens + 1) + pos] = -1; + } +} + +template +cudaError_t ChainSpeculativeSampling(DType* draft_probs, IdType* draft_token_ids, + DType* target_probs, IdType* output_token_ids, + IdType* output_accepted_token_num, + IdType* output_emitted_draft_token_num, uint32_t batch_size, + uint32_t num_speculative_tokens, uint32_t d, + bool deterministic, uint64_t philox_seed, + uint64_t philox_offset, cudaStream_t stream = 0) { + constexpr uint32_t BLOCK_THREADS = 1024; + const uint32_t vec_size = std::gcd(16 / sizeof(DType), d); + + const uint32_t smem_size = sizeof(SamplingTempStorage); + dim3 nblks(batch_size); + dim3 nthrs(BLOCK_THREADS); + void* args[] = {&draft_probs, + &draft_token_ids, + &target_probs, + &output_token_ids, + &output_accepted_token_num, + &output_emitted_draft_token_num, + &num_speculative_tokens, + &d, + &philox_seed, + &philox_offset}; + DISPATCH_ALIGNED_VEC_SIZE( + vec_size, VEC_SIZE, {DISPATCH_DETERMINISTIC(deterministic, DETERMINISTIC, { + auto kernel = ChainSpeculativeSampling; + FLASHINFER_CUDA_CALL( + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + FLASHINFER_CUDA_CALL( + cudaLaunchKernel((void*)kernel, nblks, nthrs, args, smem_size, stream)); + })}); + return cudaSuccess; +} + +} // namespace sampling + +} // namespace flashinfer + +#endif // FLASHINFER_SAMPLING_CUH_ diff --git a/include/flashinfer/semaphore_utils.cuh b/include/flashinfer/semaphore_utils.cuh new file mode 100644 index 0000000000000000000000000000000000000000..0a7485238600496ffca4f50e98bc5d67ae3c37c1 --- /dev/null +++ b/include/flashinfer/semaphore_utils.cuh @@ -0,0 +1,53 @@ +/* + * Copyright (c) 2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef FLASHINFER_SEMAPHORE_UTILS_CUH +#define FLASHINFER_SEMAPHORE_UTILS_CUH + +#include + +#include "utils.cuh" + +namespace flashinfer { + +template +__global__ void zero_gmem_semaphore(T* semaphore, int size) { + for (int i = threadIdx.x; i < size; i += blockDim.x) { + semaphore[i] = 0; + } +} + +template +cudaError_t zero_gmem_semaphore_launcher(T* semaphore, int size, bool enable_pdl, + cudaStream_t stream) { + cudaLaunchConfig_t config = {0}; + config.gridDim = 1; + config.blockDim = 128; + config.dynamicSmemBytes = 0; + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = enable_pdl; + config.numAttrs = 1; + config.attrs = attrs; + + FLASHINFER_CUDA_CALL(cudaLaunchKernelEx(&config, zero_gmem_semaphore, semaphore, size)); + + return cudaSuccess; +} + +} // namespace flashinfer + +#endif // FLASHINFER_SEMAPHORE_UTILS_CUH diff --git a/include/flashinfer/trtllm/batched_gemm/KernelRunner.h b/include/flashinfer/trtllm/batched_gemm/KernelRunner.h new file mode 100644 index 0000000000000000000000000000000000000000..4c7d296b917dbf2316141a2b8cc1b299c256f8b4 --- /dev/null +++ b/include/flashinfer/trtllm/batched_gemm/KernelRunner.h @@ -0,0 +1,129 @@ +/* + * Copyright (c) 2020-2025, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +#include +#include +#include + +#include "trtllmGen_bmm_export/Enums.h" +#include "trtllmGen_bmm_export/trtllm/gen/DtypeDecl.h" + +namespace tensorrt_llm { +namespace kernels { + +// Keep this in sync with the ActType in +// cpp/tensorrt_llm/kernels/trtllmGenKernels/batchedGemm/trtllmGen_bmm_export/GemmGatedActOptions.h +enum class ActType { + // For ActType == SwiGlu, ideally we would like to have something like + // gatedAct = scaleC * (x0 * scaleAb + beta) * ((x1 * scaleGate) * sigmoid(alpha * x1 * + // scaleGate)). + // But for now, we use the simplified version + // gatedAct = scaleC' * (x0 + beta') * ((x1 * scaleGate) * sigmoid(alpha * x1 * scaleGate)), + // where x0 and x1 are the raw numbers from Gemm, while scaleC and scaleGate are input scales, + // beta' = beta / scaleAb, scaleC' = scaleC * scaleAb. + // + // GatedSilu is a special case of SwiGlu where the alpha is 1.0 and the beta is 0.0. + SwiGlu +}; + +struct TrtllmGenBatchedGemmRunnerOptions { + batchedGemm::trtllm::gen::Dtype dtypeA; + batchedGemm::trtllm::gen::Dtype dtypeB; + batchedGemm::trtllm::gen::Dtype dtypeC; + ActType actType{ActType::SwiGlu}; + bool deepSeekFp8{false}; + bool fusedAct{false}; + bool routeAct{false}; + bool staticBatch{false}; + bool transposeMmaOutput{false}; + int32_t tileSize{8}; + int32_t epilogueTileM{128}; + bool useShuffledMatrixA{false}; + batchedGemm::gemm::MatrixLayout weightLayout{batchedGemm::gemm::MatrixLayout::MajorK}; +}; + +class TrtllmGenBatchedGemmRunner { + public: + explicit TrtllmGenBatchedGemmRunner(TrtllmGenBatchedGemmRunnerOptions const& options); + + [[nodiscard]] size_t getWorkspaceSizeInBytes(int32_t m, int32_t n, int32_t k, + std::vector const& batchedTokens, + int32_t numTokens, int32_t numBatches, + int32_t maxNumCtasInBatchDim, + int32_t configIndex) const; + + // Generic GEMM interface + void run(int32_t m, int32_t n, int32_t k, std::vector const& batchedTokens, + int32_t numTokens, int32_t numBatches, int32_t maxNumCtasInBatchDim, void const* a, + void const* sfA, void const* b, void const* sfB, void const* perTokensSfA, + void const* perTokensSfB, float const* scaleC, float const* scaleGateC, + float const* bias, float const* swiGluAlpha, float const* swiGluBeta, + float const* clampLimit, void* c, void* outSfC, int32_t const* routeMap, + int32_t const* totalNumPaddedTokens, int32_t const* ctaIdxXyToBatchIdx, + int32_t const* ctaIdxXyToMnLimit, int32_t const* numNonExitingCtas, void* workspace, + CUstream stream, int device, int32_t configIndex); + + // NVFP4 per-block scaling GEMM + void run(int32_t m, int32_t n, int32_t k, std::vector const& batchedTokens, + void const* a, void const* sfA, void const* b, void const* sfB, void* c, void* outSfC, + void* workspace, CUstream stream, int device, int32_t configIndex); + + void run(int32_t m, int32_t n, int32_t k, std::vector const& batchedTokens, + void const* a, void const* sfA, void const* b, void const* sfB, float const* bias, + float const* swiGluAlpha, float const* swiGluBeta, float const* clampLimit, void* c, + void* outSfC, void* workspace, CUstream stream, int device, int32_t configIndex); + + // FP8 per-tensor scaling GEMM + void run(int32_t m, int32_t n, int32_t k, std::vector const& batchedTokens, + void const* a, void const* b, float const* scaleC, float const* scaleGateC, void* c, + void* workspace, CUstream stream, int device, int32_t configIndex); + + // Get the list of configs that passed the validation based on the constructor options + [[nodiscard]] std::vector getPassingConfigIndices() const { + return mPassingConfigIndices; + } + + // Get the list of config indices that are valid for the given problem shape + [[nodiscard]] std::vector getValidConfigIndices( + int32_t m, int32_t n, int32_t k, std::vector const& batchedTokens, int32_t numTokens, + int32_t numBatches, int32_t maxNumCtasInBatchDim) const; + + // Get a default config index that is valid for the given problem shape + // This will be used as the fallback config if using auto-tuning + [[nodiscard]] int64_t getDefaultValidConfigIndex(int32_t m, int32_t n, int32_t k, + std::vector const& batchedTokens, + int32_t numTokens, int32_t numBatches, + int32_t maxNumCtasInBatchDim) const; + + [[nodiscard]] bool isValidConfigIndex(int32_t configIndex, int32_t m, int32_t n, int32_t k, + std::vector const& batchedTokens, + int32_t numTokens, int32_t numBatches, + int32_t maxNumCtasInBatchDim) const; + + private: + void selectGemmConfig(int32_t m, int32_t n, int32_t k, std::vector const& batchedTokens, + int32_t numTokens, int32_t numBatches, int32_t maxNumCtasInBatchDim); + + private: + TrtllmGenBatchedGemmRunnerOptions mOptions; + std::vector mPassingConfigIndices; +}; +} // namespace kernels +} // namespace tensorrt_llm diff --git a/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/BatchedGemmEnums.h b/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/BatchedGemmEnums.h new file mode 100644 index 0000000000000000000000000000000000000000..27955d2bdc4351f5e43219cb98bd21a30eecee72 --- /dev/null +++ b/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/BatchedGemmEnums.h @@ -0,0 +1,55 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & + * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include +#include + +namespace batchedGemm { + +namespace batchedGemm { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +enum class RouteImpl { + // No Routing + NoRoute = 0, + // Use LDGSTS to do the routing + Ldgsts = 1, + // Use UTMALDG.GATHER4 to do the routing + Tma = 2 +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline bool doesRouteImplUseNoRoute(RouteImpl mode) { return (mode == RouteImpl::NoRoute); } + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline bool doesRouteImplUseLdgsts(RouteImpl mode) { return (mode == RouteImpl::Ldgsts); } + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline bool doesRouteImplUseTma(RouteImpl mode) { return (mode == RouteImpl::Tma); } + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace batchedGemm + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace batchedGemm diff --git a/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/BatchedGemmInterface.h b/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/BatchedGemmInterface.h new file mode 100644 index 0000000000000000000000000000000000000000..fb81ff077a4c3414ee57b348f1c8e0254a7d84e0 --- /dev/null +++ b/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/BatchedGemmInterface.h @@ -0,0 +1,767 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & + * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include +#include + +#include "BatchedGemmOptions.h" +#include "KernelParams.h" +#include "trtllm/gen/CudaKernelLauncher.h" + +#ifdef TLLM_GEN_EXPORT_INTERFACE +#include "KernelMetaInfo.h" +#endif // TLLM_GEN_EXPORT_INTERFACE + +#ifdef TLLM_GEN_BMM_CUBIN_PATH +static const std::string tllm_gen_bmm_cubin_path = std::string(TLLM_GEN_BMM_CUBIN_PATH); +#else +static_assert(false, "TLLM_GEN_BMM_CUBIN_PATH macro is not defined when compiling"); +#endif + +namespace flashinfer::trtllm_cubin_loader { +std::string getCubin(const std::string& kernelName, const std::string& sha256); +} + +namespace batchedGemm { + +namespace batchedGemm { + +//////////////////////////////////////////////////////////////////////////////////////////////////// +// +// BatchedGemmData +// +//////////////////////////////////////////////////////////////////////////////////////////////////// + +struct BatchedGemmData { + struct ProblemDimensions { + // The number of batches. + int32_t mNumBatches{0}; + // The number of tokens. Set to 0 if routeAct is false. + int32_t mNumTokens{0}; + // Whether the batch is on the M dimension. + bool mBatchM{true}; + // The maximum number of CTAs in the token dimension. + // Need to be set if mNumTokens > 0 and the token per batch + // distribution is not known at launch time. + // In this case, the kernel will launch mMaxNumCtasInTokenDim CTAs in token dim and exit early + // if the idx of CTAs is larger or equal to mPtrNumNonExitingCtas. + int32_t mMaxNumCtasInTokenDim{0}; + + // Either mBatchedM or mBatchedN must be set when mNumTokens == 0, otherwise not used. + // The number of tokens in each batch on the M dimension if batchM, + // otherwise not used. + // The number of elements in the array is mNumBatches. + // E.g. to implement a BMM with each batch having M tokens, one needs to set mBatchedM to + // {M, M, M, .. mNumBatches times ..} + std::vector mBatchedM{}; + // The number of tokens in each batch on the N dimension if batchN, + // otherwise not used. + // The number of elements in the array is mNumBatches. + // E.g. to implement a BMM with each batch having N tokens, one needs to set mBatchedN to + // {N, N, N, .. mNumBatches times ..} + std::vector mBatchedN{}; + + // The M dimension. + // It is the total number of tokens if A is the activation matrix. + // It is the total number of output channels if A is the weight matrix. + int32_t mM{0}; + // The N dimension. + // It is the total number of tokens if B is the activation matrix. + // It is the total number of output channels if B is the weight matrix. + int32_t mN{0}; + // The K dimension. It is the hidden dimension of the input matrices. + int32_t mK{0}; + // The rank id of the current device in the multi-gpu space. + int32_t mRank{0}; + // The number of devices in tensor-parallel group. + int32_t mWorldSize{1}; + }; + + struct InputBuffers { + // The matrix A. The data type is controlled by options.mDtypeA. + // + // If (routeAct == true && batchM), the shape is [M, K] + // Else + // If batchM: + // Logical shape is [sum(divUpMul(M[bi], tileM) for bi in B), K]. + // Logical strides are [K, 1]. + // + // If batchN: + // If layoutA is MatrixLayout::MajorK + // Logical shape is [B, divUpMul(M, tileM), K]. + // Logical strides are [divUpMul(M, tileM) * K, K, 1]. + // If layoutA is MatrixLayout::MajorMn + // Logical shape is [B, K, divUpMul(M, tileM)]. + // Logical strides are [K * divUpMul(M, tileM), divUpMul(M, tileM), 1]. + // If layoutA is MatrixLayout::BlockMajorK + // Logical shape is [B, K / blockK, divUpMul(M, tileM), blockK]. + // Logical strides are [K * divUpMul(M, tileM), divUpMul(M, tileM) * blockK, blockK, 1]. + // where blockK is 128B. + void const* mPtrA{nullptr}; + + // The block scaling factors to dequantize A. + // + // If (routeAct == true && batchM), the shape is [M, K / 16] + // Else + // If DeepSeek FP8 recipe is used: + // If transposeMmaOutput is false, shape is [K / 128, M]. + // Otherwise, shape is [M / 128, K / 128]. + // The rightmost dimension is contiguous in memory. + // + // If DeepSeek FP8 recipe is not used, but for MxFp{4,8} and NvFp4 formats: + // The layout of scaling factors for A is always R128c4 + // M must be a multiple of 128. + // K must be a multiple of 64. + // The "logical" shape is: [paddedM, K / 16]. + // The R128c4 layout is: [paddedM / 128, K / 16 / 4, 512]. + // The shape we use for TMA is: [paddedM / 128, K / 16 / 4, 2, 256]. + // Where paddedM is M if (routeAct == true && batchM), or + // sum(divUpMul(M[bi], tileM) for bi in B) if batchM, + // otherwise divUpMul(M, tileM) * B. + // Dtype is Dtype::Fp32 if DeepSeek FP8 recipe is used, otherwise Dtype::E4m3. + // + // Otherwise should be set to nullptr. + void const* mPtrSfA{nullptr}; + + // The per-token scaling factors from scale A. + // + // This is used for either: + // * Per-token scaling factor quantization schemes, such as MetaFP8. The dtype is + // Dtype::Float32 + // * When the routing scales are applied to the input activations (only when output is not + // transposed). The dtype is Dtype::Bfloat16 + // + // if (batchM (A is activations)): + // Logical shape is [sum(divUpMul(M[bi], tileM) for bi in B)] + // + // if (batchN (A is weights)): + // Logical shape is [B, divUpMul(M, tileM)] + // + void const* mPtrPerTokenSfA{nullptr}; + + // The matrix B. The data type is controlled by options.mDtypeB. + // + // If (routeAct == true && batchN), the shape is [N, K] + // + // Else + // If batchN: + // Logical shape is [sum(divUpMul(N[bi], tileN) for bi in B), K]. + // Logical strides are [K, 1]. + // + // If batchM: + // If layoutB is MatrixLayout::MajorK + // Logical shape is [B, divUpMul(N, tileN), K]. + // Logical strides are [divUpMul(N, tileN) * K, K, 1]. + // If layoutB is MatrixLayout::MajorMn + // Logical shape is [B, K, divUpMul(N, tileN)]. + // Logical strides are [K * divUpMul(N, tileN), divUpMul(N, tileN), 1]. + // If layoutB is MatrixLayout::BlockMajorK + // Logical shape is [B, K / blockK, divUpMul(N, tileN), blockK]. + // Logical strides are [K * divUpMul(N, tileN), divUpMul(N, tileN) * blockK, blockK, 1]. + // where blockK is 128B. + void const* mPtrB{nullptr}; + + // The scaling factors to dequantize B. + // + // + // + // Else + // If DeepSeek FP8 recipe is used: + // If transposeMmaOutput is false, shape is [paddedN / 128, K / 128]. + // Otherwise, shape is [K / 128, paddedN]. + // The rightmost dimension is contiguous in memory. + // + // If DeepSeek FP8 recipe is not used, but for MxFp{4,8} and NvFp4 formats: + // If the layout is R128c4, + // paddedN must be a multiple of 128. + // K must be a multiple of 64. + // The R128c4 layout is: [paddedN / 128, K / 16 / 4, 512] + // The shape we use for TMA is: [paddedN / 128, K / 16 / 4, 2, 256] + // + // If the layout is R8c4, + // paddedN must be a multiple of 8. + // K must be a multiple of 64. + // The R8c4 layout is: [paddedN / 8, K / 16 / 4, 32] + // The shape we use for TMA is: [paddedN / 8, K / 16 / 4 / repeats, repeats * 32] + // where repeats = min(tileK / 16 / 4, 8) + // + // where paddedN is N if (routeAct == true && batchN), + // or sum(divUpMul(N[bi], tileN) for bi in B) if batchN, + // otherwise divUpMul(N, TileN) * B. + // + // Dtype is Dtype::Fp32 if DeepSeek FP8 recipe is used, otherwise Dtype::E4m3. + // + // Otherwise should be set to nullptr. + void const* mPtrSfB{nullptr}; + + // The per-token scaling factors from scale B. + // + // This is used for either: + // * Per-token scaling factor quantization schemes, such as MetaFP8. The dtype is + // Dtype::Float32 + // * When the routing scales are applied to the input activations (only when output is + // transposed). The dtype is Dtype::Bfloat16 + // + // if (batchM (B is weights)): + // Logical shape is [B, divUpMul(N, tileN)] + // + // if (batchN (B is activations)): + // Logical shape is [sum(divUpMul(N[bi], tileN) for bi in B)] + void const* mPtrPerTokenSfB{nullptr}; + + // The bias applied after the GEMM and before the activation function. + // The bias is applied before applying the global scaling factor. I.e. + // C = act(A * B + bias') * scaleC + // scaleC = dequantA * dequantB * quantC + // Thus, the bias' = bias / (dequantA * dequantB), where the bias is the original bias. + // + // If batchM, BiasType must be N, and bias shape is [B, N]. + // The bias is broadcasted along the M dimension. + // + // If batchN BiasType must be M, and bias shape is [B, M]. + // The bias is broadcasted along the N dimension. + // + // The dtype is float32. + void const* mPtrBias{nullptr}; + + // The output tensor scaling factor for MxFp{4,8}, Fp8 and NvFp4 quantization. + // TensorRT-LLM API requires a scaling factor on the device. + // Shape is [B]. + float const* mPtrScaleC{nullptr}; + + // The output gate scale for MxFp{4,8} and NvFp4 quantization. + // TensorRT-LLM API requires a scaling factor on the device. + // Shape is [B]. + float const* mPtrScaleGate{nullptr}; + + // The clamp limit for the accumulator before applying the activation. + // Shape is [B]. + // Clamp is INF if nullptr. + // When the input is FP8 or NVFP4, the clamp has to be scaled by limit' = limit / dequantAb. + // If applied on SwiGlu, it will be: + // + // x_glu = x_glu.clamp(min=None, max=limit) + // x_linear = x_linear.clamp(min=-limit, max=limit) + // + // The given clamp limit applies to the dequantized values, so the order of operations would + // look something like this: + // + // x0 = x0 * dqAb + // x0 = clamp(x0, none, limit) + // x0 = x0 * sigmoid(alpha * x0) + // x1 = dqAb * x1 + // x1 = clamp(x1, -limit, limit) + // out = qC * (x1 + beta) * x0 + // + // Given that the dqAb and qC are combined into scaleC, we can bring the dqAb into the clamp + // limit and apply the clamping prior to dequantization: + // + // x0 = clamp(x0, none, limit / dqAb) + // x0 = x0 * dqAb + // x0 = x0 * sigmoid(alpha * x0) + // x1 = clamp(x1, -limit / dqAb, limit / dqAb) + // scaleC = dqAb * qC + // beta' = beta / dqAb + // out = scaleC * (x1 + beta') * x0 + // + // Note this assumes that scaleAb == scaleGate which is true in TRT-LLM MoE use-case + // + float const* mPtrClampLimit{nullptr}; + + // The alpha and beta for SwiGlu. + // gatedActivation <- (x0 + beta) * activation(x1, alpha) + // Shape is [B]. + // Alpha is 1.f if nullptr. + // Beta is 0.f if nullptr. + // The formula: + // + // out_glu = x_glu * torch.sigmoid(alpha * x_glu) + (x_linear + beta) + float const* mPtrSwiGluAlpha{nullptr}; + float const* mPtrSwiGluBeta{nullptr}; + + // Param is used when the kernel is configured with -routeAct true. + // The inputs are not padded, but the outputs are padded to divUpMul(M[bi], tileM) for batchM or + // divUpMul(N[bi], tileN) for batchN. + // If -routeAct is false, the params are not used and should be set to zero. + + // The routeMap for the input tokens. + // Map of expanded token index (counting the previous padded tokens) to the batch index + // the token belongs to. + // The shape is + // [divUpMul(numTokens + numBatches * (tileM/N - 1), tileM/N)] + // The dtype is int32_t. + // + // There are 3 tokens [0, 1, 2] such that [0, 1] belong to batch [B0] and [2] to batch [B1]. + // Let's assume that the padded size is 4. + // + // The expanded indices for tokens [0, 1, 2] are: + // expandedIdx[0] = 0 + // expandedIdx[1] = 1 + // expandedIdx[2] = divUpMul(2, 4) + 0 = 4 + // + // The route map is [B0, B0, X, X, B1, X, X, X] where X could be any value. + int32_t const* mPtrRouteMap; + + ////////////////////////////////////////////////////////////////////////////////////////////////// + // + // Batching information parameters. + // + ////////////////////////////////////////////////////////////////////////////////////////////////// + + // In some cases, some CTAs must early-exit. E.g. when the grid size is set statically, but the + // actual workload is decided at runtime. This element on the device contains the number of CTAs + // that do not early-exit. The number corresponds to the X dim of the grid when the output is + // not transposed (i.e. batchM). To the Y dim, otherwise. The size is 1 and the dtype is + // int32_t. Used if isStaticBatch == false, otherwise set to nullptr. The pointer points to a + // scalar and the dtype is int32_t. The pointed value must be >= 0. + int32_t const* mPtrNumNonExitingCtas; + + // Pointer to total number of padded tokens. + // Computed as + // int32_t totalNumPaddedTokens{0}; + // for (int bi = 0; bi < options.mNumBatches; bi++) { + // totalNumPaddedTokens += batchM ? divUpMul(options.mBatchedM[bi], options.mTileM) + // : divUpMul(options.mBatchedN[bi], options.mTileN); + // } + // The size is 1 and the dtype is int32_t. + // If isStaticBatch == true, ptrTotalNumPaddedTokens should be set to nullptr and + // totalNumPaddedTokens is used. + int32_t const* mPtrTotalNumPaddedTokens; + + // Pointer to the map from the CTA index (in X/Y dim) to the batch index. + // Maps CTA index in batch dim (i.e. blockDim.x if batchM, otherwise blockDim.y) + // to batch index. + // E.g. with listM = 128,255,32 and tileM = 128, should be equal to + // ctaIdxXyToBatchIdx = [0, 1, 1, 2] + // If isStaticBatch == true, ptrCtaIdxXyToBatchIdx should be set to nullptr and + // ctaIdxXyToBatchIdx is used. + // The shape is + // [divUp(numTokens + numBatches * (tileM/N - 1), tileM/N)] + int32_t const* mPtrCtaIdxXyToBatchIdx; + + // Pointer from the CTA index X/Y to the expanded tile index where the expanded tile index is + // computed as: + // + // int expandedIdx = 0; + // for (int bi = 0; bi < batchIdx-1; ++bi) { + // expandIdx = divUpMul(numTokens[bi], TileM/N); + // } + // expandIdx += + // E.g. with numTokens = [128,255,32] and tileM = 128, should be equal to + // ptrCtaIdxXyToMnLimit = [128, 256, 383, 416] + // The shape is + // [divUp(numTokens + numBatches * (tileM/N - 1), tileM/N)] + int32_t const* mPtrCtaIdxXyToMnLimit; + }; + + struct OutputBuffers { + // The output matrix C. The data type is controlled by options.mDtypeC. + // + // If batchM: + // Logical shape is [sum(divUpMul(M[bi], tileM) for bi in B), N]. + // Logical strides are [N, 1]. + // + // If batchN: + // Logical shape is [sum(divUpMul(N[bi], tileN) for bi in B), M]. + // Logical strides are [M, 1]. + void* mPtrC{nullptr}; + + // The output block scaling factors for C. + // + // If MxFp{4,8} and NvFp4 formats are used, + // The "logical" shape is: + // if batchM: [paddedM, N / 16] + // if batchN: [paddedN, M / 16] + // where paddedM is sum(divUpMul(M[bi], tileM) for bi in B), + // where paddedN is sum(divUpMul(N[bi], tileN) for bi in B). + // + // If the layout is R128c4, + // paddedOuter must be a multiple of 128. + // inner must be a multiple of 64. + // The R128c4 layout is: [paddedOuter / 128, inner / 16 / 4, 512] + // The shape we use for TMA is: [paddedOuter / 128, inner / 16 / 4, 2, 256] + // where inner = N if batchM, otherwise M. + // where paddedOuter = paddedM if batchM, otherwise paddedN. + // + // If the layout is R8c4, + // paddedOuter must be a multiple of 8. + // inner must be a multiple of 64. + // The R8c4 layout is: [paddedOuter / 8, inner / 16 / 4, 32] + // The shape we use for TMA is: [paddedOuter / 8, inner / 16 / 4 / repeats, repeats * 32] + // where repeats = min(tileInner / 16 / 4, 8), + // where tileInner = tileN if batchM, otherwise tileM, + // where paddedOuter = paddedM if batchM, otherwise paddedN. + // where inner = N if batchM, otherwise M. + // + // The dtype is Dtype::E4m3. + // + // If DeepSeek FP8 quantization recipe is used, + // If batchM: + // The shape is [N / 128, paddedM], + // where paddedM is sum(divUpMul(M[bi], tileM) for bi in B). + // If batchN: + // The shape is [M / 128, paddedN], + // where paddedN is sum(divUpMul(N[bi], tileN) for bi in B). + // The rightmost dimension is contiguous in memory. + // The dtype is Dtype::Float32. + void* mPtrSfC{nullptr}; + }; + + ProblemDimensions mProblemDimensions; + InputBuffers mInputBuffers; + OutputBuffers mOutputBuffers; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// +// +// BatchedGemmInterface +// +//////////////////////////////////////////////////////////////////////////////////////////////////// + +class BatchedGemmInterface { + public: + using ModuleCache = std::unordered_map>; + + BatchedGemmInterface() {} + + // Launch the cubin from the provided config. It calls all necessary memsets for internal buffers. + // Provided config must be validated with isValidConfig before the call. + int32_t run(BatchedGemmConfig const& config, void* workspace, BatchedGemmData const& options, + void* cudaStream, int32_t multiProcessorCount, bool usePdl = true, + std::optional> moduleCache = std::nullopt); + + // Initializes the buffers before the world sync. Must be called before run. + int32_t runInitBeforeWorldSync(BatchedGemmConfig const& /* config */, + BatchedGemmData const& /* data */, void* /* cudaStream */) const { + return 0; + }; + + size_t getWorkspaceSizeInBytes(BatchedGemmConfig const& /* config */, + BatchedGemmData const& /* data */) const; + + // Returns the list of all available cubin configurations + BatchedGemmConfig const* getBatchedGemmConfigs() const; + + // Returns the number of available cubin configurations + size_t getNumBatchedGemmConfigs() const; + + // Returns the number of CTAs of the last launched kernel. + int32_t getNumCtas() const { return mNumCtas; } + + // Returns true if the configuration of the cubin can be executed for the given params. + bool isValidConfig(BatchedGemmConfig const& config, BatchedGemmData const& data) const; + + private: + // Aligns the pointer to the alignment + template + inline Dtype* alignPtr(Dtype* ptr, int64_t alignment) const; + // Creates GemmOptions from kernel and data. + BatchedGemmOptions getOptionsFromConfigAndData(BatchedGemmConfig const& config, + BatchedGemmData const& data) const; + + // Returns the size of the workspace buffers in bytes + std::vector getWorkspaceSizesInBytes(BatchedGemmConfig const& config, + BatchedGemmData const& data) const; + + // Returns the size padded to the alignment + size_t getSizePaddedToAlignment(size_t size, size_t alignment) const; + + private: + // Number of the CTAs of the last launched kernel. + int32_t mNumCtas{0}; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +inline Dtype* BatchedGemmInterface::alignPtr(Dtype* ptr, int64_t alignment) const { + assert((alignment & (alignment - 1)) == 0 && "Alignment must be a power of 2"); + return reinterpret_cast((reinterpret_cast(ptr) + alignment - 1) & + ~(alignment - 1)); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +BatchedGemmConfig const* BatchedGemmInterface::getBatchedGemmConfigs() const { +#ifdef TLLM_GEN_EXPORT_INTERFACE + return tensorrt_llm::kernels::tllmGenBatchedGemmList; +#else + return nullptr; +#endif +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +size_t BatchedGemmInterface::getNumBatchedGemmConfigs() const { +#ifdef TLLM_GEN_EXPORT_INTERFACE + return tensorrt_llm::kernels::tllmGenBatchedGemmListLen; +#else + return 0; +#endif +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +BatchedGemmOptions BatchedGemmInterface::getOptionsFromConfigAndData( + BatchedGemmConfig const& config, BatchedGemmData const& data) const { + // Create options from config and data. + BatchedGemmOptions options; + options = config.mOptions; + options.mM = data.mProblemDimensions.mM; + options.mN = data.mProblemDimensions.mN; + options.mK = data.mProblemDimensions.mK; + options.mBatchedM = data.mProblemDimensions.mBatchedM; + options.mBatchedN = data.mProblemDimensions.mBatchedN; + options.mBatchMode = data.mProblemDimensions.mBatchM ? BatchedGemmOptions::BatchMode::BatchM + : BatchedGemmOptions::BatchMode::BatchN; + options.mNumBatches = data.mProblemDimensions.mNumBatches; + options.mNumTokens = data.mProblemDimensions.mNumTokens; + return options; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +bool BatchedGemmInterface::isValidConfig(BatchedGemmConfig const& config, + BatchedGemmData const& data) const { + // Get options from config and data. + auto options = getOptionsFromConfigAndData(config, data); + + // Is Blackwell? + bool isBlackwell = config.mSm == gemm::SmVersion::Sm100a; + + // Check options without modifications. + return checkAndUpdateBatchedGemmOptions(options, isBlackwell, + /* updateOptions */ false); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +size_t BatchedGemmInterface::getSizePaddedToAlignment(size_t size, size_t alignment) const { + assert((alignment & (alignment - 1)) == 0); + return (size + alignment - 1) & ~(alignment - 1); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +size_t BatchedGemmInterface::getWorkspaceSizeInBytes(BatchedGemmConfig const& config, + BatchedGemmData const& data) const { + auto workspaceSizes = getWorkspaceSizesInBytes(config, data); + auto size = std::accumulate(workspaceSizes.begin(), workspaceSizes.end(), 0); + // Additional 1023 bytes to align the pointer to 1024 + return size > 0 ? size + 1023 : 0; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +std::vector BatchedGemmInterface::getWorkspaceSizesInBytes( + BatchedGemmConfig const& config, BatchedGemmData const& data) const { + std::vector workspaceSizes; + + // Get options from config and data. + auto options = getOptionsFromConfigAndData(config, data); + + if (options.mUseDeepSeekFp8 && options.mFusedAct) { + int32_t totalNumPaddedTokens = 0; + auto const batchM = options.mBatchMode == BatchedGemmOptions::BatchMode::BatchM; + if (!options.mEnablesEarlyExit || options.mNumTokens == 0) { + for (int32_t bi = 0; bi < options.mNumBatches; ++bi) { + totalNumPaddedTokens += batchM ? gemm::divUpMul(options.mBatchedM[bi], options.mTileM) + : gemm::divUpMul(options.mBatchedN[bi], options.mTileN); + } + } else { + // Get tile in token dim. + auto tileTokensDim = batchM ? options.mTileM : options.mTileN; + totalNumPaddedTokens = data.mProblemDimensions.mMaxNumCtasInTokenDim * tileTokensDim; + } + + // Get options from config. + auto& options = config.mOptions; + + int const tokenTile = batchM ? options.mTileM : options.mTileN; + + auto const numTokens = totalNumPaddedTokens; + auto const intermediateDim = batchM ? options.mN : options.mM; + auto const intermediateTile = batchM ? options.mTileN : options.mTileM; + + auto const numBytesRowMax = intermediateDim * totalNumPaddedTokens / 128 * sizeof(float); + + auto const numTilesToken = numTokens / tokenTile; + auto const numTilesInt = intermediateDim / intermediateTile; + auto const numBytesRowMaxBars = numTilesToken * numTilesInt / 2 * sizeof(uint32_t); + + // TODO: do we need to pad to 1024? + workspaceSizes.push_back(getSizePaddedToAlignment(numBytesRowMax, 1024)); + workspaceSizes.push_back(getSizePaddedToAlignment(numBytesRowMaxBars, 1024)); + } + + return workspaceSizes; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// +int32_t BatchedGemmInterface::run(BatchedGemmConfig const& config, void* workspace, + BatchedGemmData const& batchedGemmData, void* cudaStream, + int32_t /* multiProcessorCount */, bool usePdl, + std::optional> moduleCache) { + // Might be used. + (void)usePdl; + (void)moduleCache; + // Get options from config and data. + auto options = getOptionsFromConfigAndData(config, batchedGemmData); + + bool const batchM = options.mBatchMode == BatchedGemmOptions::BatchMode::BatchM; + bool const useDeepSeekFp8 = options.mUseDeepSeekFp8 && options.mDtypeA == tg::Dtype::E4m3 && + options.mDtypeB == tg::Dtype::E4m3; + + auto workspaceSizes = getWorkspaceSizesInBytes(config, batchedGemmData); + float* dPtrRowMax{nullptr}; + uint32_t* dPtrRowMaxBars{nullptr}; + + // Set the completion barriers to 0 if needed. + if (useDeepSeekFp8 && options.mFusedAct) { + dPtrRowMax = reinterpret_cast(alignPtr(reinterpret_cast(workspace), 1024)); + dPtrRowMaxBars = reinterpret_cast( + alignPtr(reinterpret_cast(dPtrRowMax) + workspaceSizes[0], 1024)); + auto err = cudaMemsetAsync((void*)dPtrRowMaxBars, 0x00, workspaceSizes[1], + reinterpret_cast(cudaStream)); + if (err != cudaSuccess) { + return 1; + } + } + + int32_t numCtaXy{0}; + if (options.mIsStaticBatch) { + for (int32_t bi = 0; bi < options.mNumBatches; ++bi) { + numCtaXy += batchM ? gemm::divUp(options.mBatchedM[bi], options.mTileM) + : gemm::divUp(options.mBatchedN[bi], options.mTileN); + } + } + + int32_t maxNumCtasInBatchDim{numCtaXy}; + // For normal BMM, mNumTokens == 0 and the number of CTAs is known to host. + // For MoE, mNumTokens != 0 and the number of CTAs is known only at runtime. + // We launch maximally possible number of CTAs and use ptrNumNonExitingCtas to determine + // the actual number of CTAs to run. + if ((options.mEnablesEarlyExit || options.mEnablesDelayedEarlyExit) && options.mNumTokens != 0) { + // Get maximum number of CTAs in batch dim. + maxNumCtasInBatchDim = batchedGemmData.mProblemDimensions.mMaxNumCtasInTokenDim; + } + + auto const numCtaX = batchM ? maxNumCtasInBatchDim : gemm::divUp(options.mM, options.mTileM); + auto const numCtaY = batchM ? gemm::divUp(options.mN, options.mTileN) : maxNumCtasInBatchDim; + auto const numCtaZ = options.mNumSlicesForSplitK; + mNumCtas = numCtaX * numCtaY * numCtaZ; + + auto kernelParams = KernelParamsSetup::setKernelParams( + options, batchM, batchedGemmData.mInputBuffers.mPtrA, batchedGemmData.mInputBuffers.mPtrB, + batchedGemmData.mOutputBuffers.mPtrC, batchedGemmData.mInputBuffers.mPtrSfA, + batchedGemmData.mInputBuffers.mPtrSfB, batchedGemmData.mInputBuffers.mPtrPerTokenSfA, + batchedGemmData.mInputBuffers.mPtrPerTokenSfB, batchedGemmData.mInputBuffers.mPtrBias, + batchedGemmData.mOutputBuffers.mPtrSfC, batchedGemmData.mInputBuffers.mPtrScaleC, + batchedGemmData.mInputBuffers.mPtrScaleGate, batchedGemmData.mInputBuffers.mPtrClampLimit, + batchedGemmData.mInputBuffers.mPtrSwiGluAlpha, batchedGemmData.mInputBuffers.mPtrSwiGluBeta, + batchedGemmData.mInputBuffers.mPtrRouteMap, dPtrRowMax, dPtrRowMaxBars, + batchedGemmData.mInputBuffers.mPtrNumNonExitingCtas, + batchedGemmData.mInputBuffers.mPtrTotalNumPaddedTokens, + batchedGemmData.mInputBuffers.mPtrCtaIdxXyToBatchIdx, + batchedGemmData.mInputBuffers.mPtrCtaIdxXyToMnLimit, maxNumCtasInBatchDim); + + // The size of the grid. + std::vector grid{numCtaX, numCtaY, numCtaZ}; + +#ifdef TLLM_GEN_EXPORT_INTERFACE + CUmodule cuModule; + CUfunction cuFunction; + + auto fiModuleLoadData = [&](CUmodule* module) { + const std::string sha256 = config.mHash ? config.mHash : ""; + std::string fname_cubin = config.mFunctionName; + if (!fname_cubin.empty()) { + fname_cubin[0] = static_cast(std::toupper(static_cast(fname_cubin[0]))); + } + fname_cubin = tllm_gen_bmm_cubin_path + fname_cubin; + std::string cubin = flashinfer::trtllm_cubin_loader::getCubin(fname_cubin, sha256); + cuModuleLoadData(&cuModule, cubin.c_str()); + }; + + if (moduleCache.has_value()) { + ModuleCache& moduleCacheRef = moduleCache.value().get(); + + // Modules are associated with a specific context, so the context is included in the key + CUcontext ctx; + unsigned long long ctxId; + cuCtxGetCurrent(&ctx); + cuCtxGetId(ctx, &ctxId); + + // Reinterpret the ctxId as a string to avoid needing a custom hash or converting it to a + // string in decimal representation. + std::string const ctxName = + std::string(reinterpret_cast(&ctxId), sizeof(unsigned long long) / sizeof(char)); + std::string const funcName = std::string(config.mFunctionName); + auto const moduleKey = ctxName + funcName; + auto module = moduleCacheRef.find(moduleKey); + + // Use cache if module is found, otherwise load and insert into cache + if (module != moduleCacheRef.end()) { + cuFunction = std::get<1>(module->second); + } else { + fiModuleLoadData(&cuModule); + cuModuleGetFunction(&cuFunction, cuModule, config.mFunctionName); + moduleCacheRef.insert(std::make_pair(moduleKey, std::make_tuple(cuModule, cuFunction))); + } + } else { + fiModuleLoadData(&cuModule); + cuModuleGetFunction(&cuFunction, cuModule, config.mFunctionName); + } + + // Prepare the grid/block. + dim3 block3{static_cast(config.mNumThreadsPerCTA), static_cast(1), + static_cast(1)}; + dim3 grid3{(grid.size() > 0 ? static_cast(grid[0]) : 1u), + (grid.size() > 1 ? static_cast(grid[1]) : 1u), + (grid.size() > 2 ? static_cast(grid[2]) : 1u)}; + // Prepare the cluster size. + dim3 cluster3{static_cast(options.mClusterDimX), + static_cast(options.mClusterDimY), + static_cast(options.mClusterDimZ)}; + + // Run the kernel. + auto result = trtllm::gen::launchKernel( + (void*)&kernelParams, cudaStream, config.mSharedMemSize, cuFunction, block3, grid3, cluster3, + usePdl && (config.mOptions.mGridWaitForPrimaryEarlyExit | + config.mOptions.mGridWaitForPrimaryA | config.mOptions.mGridWaitForPrimaryB)); + if (result != CUDA_SUCCESS) { + return -1; + } + // If a module cache has not been given, unload the module to avoid leaking + if (!moduleCache.has_value()) { + cuModuleUnload(cuModule); + } +#else + config.mCudaRunner->run((void*)&kernelParams, (void*)cudaStream, grid); +#endif + + return 0; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace batchedGemm + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace batchedGemm diff --git a/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/BatchedGemmOptions.h b/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/BatchedGemmOptions.h new file mode 100644 index 0000000000000000000000000000000000000000..42d4d9ce1e237e2d2780ff582b8d73f4eaf9aaf9 --- /dev/null +++ b/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/BatchedGemmOptions.h @@ -0,0 +1,372 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & + * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include + +#include "BatchedGemmEnums.h" +#include "GemmGatedActOptions.h" +#include "GemmOptions.h" + +#ifndef TLLM_GEN_EXPORT_INTERFACE +#include "trtllm/gen/CudaRunner.h" +#include "trtllm/gen/GenCtx.h" +#else +#include + +#define TLLM_CHECK_ERROR(cond, ...) \ + if (!(cond)) { \ + printArgs(__VA_ARGS__); \ + return false; \ + } + +#define TLLM_LOG_ERROR(...) TLLM_CHECK_ERROR(false, __VA_ARGS__) + +#define TLLM_CHECK_ERROR_FMT(...) TLLM_CHECK_ERROR(false, __VA_ARGS__) + +#define TLLM_CHECK_WARNING(cond, ...) \ + if (!(cond)) { \ + printArgs(__VA_ARGS__); \ + return false; \ + } + +#define TLLM_LOG_WARNING(...) TLLM_CHECK_WARNING(false, __VA_ARGS__) + +#define TLLM_LOG_INFO(...) TLLM_CHECK_WARNING(false, __VA_ARGS__) + +#endif + +namespace batchedGemm { + +namespace batchedGemm { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace tg = trtllm::gen; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// We do not differentiate between BatchedGemmOptions and BatchedGemmGatedActOptions for simplicity. +// We inherit from GemmGatedActOptions, which is inherited from +// GemmOptions to get GemmOptions and GemmGatedActOptions at the same time. +struct BatchedGemmOptions : public gemmGatedAct::GemmGatedActOptions { + // Dtor. Allow down-casting. + virtual ~BatchedGemmOptions() = default; + + enum class BatchMode { BatchM, BatchN }; + + BatchedGemmOptions() = default; + // FIXME We create explicit constructor with all options to WAR stubgen issue in TRT-LLM. + BatchedGemmOptions( + gemm::AllReduceAlgo allReduceAlgo, gemm::BiasType biasType, int blockK, int clusterDimX, + int clusterDimY, int clusterDimZ, tg::Dtype dtypeAcc, tg::Dtype dtypeA, tg::Dtype dtypeB, + tg::Dtype dtypeC, tg::Dtype dtypeMmaA, tg::Dtype dtypeMmaB, bool enablesEarlyExit, + bool enablesDelayedEarlyExit, bool enablesGlobalPtxKnobs, int epilogueLdtmDps, + int epilogueLdtmBits, int epilogueTileM, int epilogueTileN, bool gridTriggerSecondaryA, + bool gridTriggerSecondaryB, bool gridWaitForPrimaryEarlyExit, bool gridWaitForPrimaryA, + bool gridWaitForPrimaryB, bool hoistLoadTaskInit, bool hoistMmaTaskTryWaits, int k, + gemm::KernelTraits kernelTraits, gemm::MatrixLayout layoutA, gemm::MatrixLayout layoutB, + int m, int mmaK, tg::MmaKind mmaKind, int mmaM, int mmaN, bool mockAllReduce, int n, + int numSlicesForSplitK, int numSlicesForSliceK, int numStages, int numStagesMma, + int numStagesMmaWithinWorkTile, int numStagesMmaAcrossWorkTile, int numStagesWorkId, + bool outputDebugTensors, bool patchF2fp, bool useShuffledMatrixA, bool sliceK, + gemm::SplitK splitK, bool transposeMmaOutput, int tileM, int tileN, int tileK, + bool useUnrollLoop2xForMma, bool useCustomMmaSchedule, + bool useHoistTryWaitForCustomMmaSchedule, bool useDeepSeekFp8, bool usePerTokenSfA, + bool usePerTokenSfB, bool useTmaStore, bool useTwoTmaLoadWarps, bool useTwoMmaWarps, + tg::SfLayout sfLayoutA, tg::SfLayout sfLayoutB, tg::SfLayout sfLayoutC, + int32_t sfReshapeFactor, gemm::TileScheduler tileScheduler, gemmGatedAct::ActType actType, + bool clampBeforeAct, std::vector batchedM, std::vector batchedN, + BatchMode batchMode, int numBatches, bool isStaticBatch, int numTokens, RouteImpl routeImpl, + bool gridWaitForPrimaryRouting, bool fusedAct, int numRegsPerThreadNonEpilogueWarp, + int numRegsPerThreadEpilogueWarp, int numRegsCastAWarps, bool useTmaOobOpt) + : gemmGatedAct::GemmGatedActOptions( + gemm::GemmOptions( + allReduceAlgo, biasType, blockK, clusterDimX, clusterDimY, clusterDimZ, dtypeAcc, + dtypeA, dtypeB, dtypeC, dtypeMmaA, dtypeMmaB, enablesEarlyExit, + enablesDelayedEarlyExit, enablesGlobalPtxKnobs, epilogueLdtmDps, epilogueLdtmBits, + epilogueTileM, epilogueTileN, gridTriggerSecondaryA, gridTriggerSecondaryB, + gridWaitForPrimaryEarlyExit, gridWaitForPrimaryA, gridWaitForPrimaryB, + hoistLoadTaskInit, hoistMmaTaskTryWaits, k, kernelTraits, layoutA, layoutB, m, mmaK, + mmaKind, mmaM, mmaN, mockAllReduce, n, numSlicesForSplitK, numSlicesForSliceK, + numStages, numStagesMma, numStagesMmaWithinWorkTile, numStagesMmaAcrossWorkTile, + numStagesWorkId, outputDebugTensors, patchF2fp, useShuffledMatrixA, sliceK, splitK, + transposeMmaOutput, tileM, tileN, tileK, useUnrollLoop2xForMma, + useCustomMmaSchedule, useHoistTryWaitForCustomMmaSchedule, useDeepSeekFp8, + usePerTokenSfA, usePerTokenSfB, useTmaStore, useTwoTmaLoadWarps, useTwoMmaWarps, + sfLayoutA, sfLayoutB, sfLayoutC, sfReshapeFactor, tileScheduler), + actType, clampBeforeAct), + mBatchedM(batchedM), + mBatchedN(batchedN), + mBatchMode(BatchMode(batchMode)), + mFusedAct(fusedAct), + mGridWaitForPrimaryRouting(gridWaitForPrimaryRouting), + mIsStaticBatch(isStaticBatch), + mNumBatches(numBatches), + mNumRegsPerThreadNonEpilogueWarp(numRegsPerThreadNonEpilogueWarp), + mNumRegsPerThreadEpilogueWarp(numRegsPerThreadEpilogueWarp), + mNumRegsCastAWarps(numRegsCastAWarps), + mNumTokens(numTokens), + mRouteImpl(routeImpl), + mUseTmaOobOpt(useTmaOobOpt) {} + + // Batched M-dimensions of GEMM. + std::vector mBatchedM; + // Batched N-dimensions of GEMM. + std::vector mBatchedN; + // Whether batching M or N. + BatchMode mBatchMode{BatchMode::BatchM}; + // Whether to perform a fused gated activation. + bool mFusedAct{false}; + // Whether the loads that load from ptrRouteMap, ptrTotalNumPaddedTokens, + // ptrCtaIdxXyToBatchIdx, etc.. should wait on a grid dependency. + bool mGridWaitForPrimaryRouting{true}; + // Whether the batch size is static (i.e. known at kernel launch time). + bool mIsStaticBatch{true}; + // Number of Gemm batches. + int mNumBatches; + // Number of registers per thread for non-epilogue warps + int mNumRegsPerThreadNonEpilogueWarp{0}; + // Number of registers per thread for epilogue warps + int mNumRegsPerThreadEpilogueWarp{0}; + // Number of registers for the cast A warps. + int mNumRegsCastAWarps{0}; + // Total number of tokens. + int mNumTokens{32}; + // Whether load the input tokens and do routing. + RouteImpl mRouteImpl{RouteImpl::NoRoute}; + // Whether to use TMA out-of-bounds optimization to reduce wasted traffic. See details in + // BatchedGemm/KernelParamsDecl.h. + bool mUseTmaOobOpt{false}; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// Check if the options are valid or not. +bool checkAndUpdateBatchedGemmOptions(BatchedGemmOptions& options, bool isBlackwell, + bool updateOptions = true) { + bool isValid = true; + if (options.mUseTmaOobOpt && !options.mUseTwoTmaLoadWarps) { + if (updateOptions) { + // Since any routing (mRouteAct != NoRoute) requires mUseTwoTmaLoadWarps == true. + // Single TMA load warp is not the target use case for OOB optimization. + options.mUseTmaOobOpt = false; + } else { + TLLM_CHECK_ERROR(false, "TMA OOB optimization requires two TMA load warps."); + return false; + } + } + if (options.mFusedAct) { + // ensure that we check the fused options as well + isValid = gemmGatedAct::checkAndUpdateGemmGatedActOptions(options, isBlackwell, updateOptions); + } else { + isValid = + gemm::checkAndUpdateGemmOptions(options, isBlackwell, 1 /* tpGrpSize */, updateOptions); + } + + bool batchM = options.mBatchMode == BatchedGemmOptions::BatchMode::BatchM; + if (updateOptions) { + if (batchM) { + if (options.mBatchedM.empty()) { + options.mBatchedM.push_back(128); + options.mBatchedM.push_back(256); + } + options.mNumBatches = options.mBatchedM.size(); + } else { + if (options.mBatchedN.empty()) { + options.mBatchedN.push_back(128); + options.mBatchedN.push_back(256); + } + options.mNumBatches = options.mBatchedN.size(); + } + } + + for (int b = 0; b < options.mNumBatches; b++) { + if (batchM) { + TLLM_CHECK_ERROR(options.mN > 0 && options.mK > 0, "N and K must be larger than 0"); + TLLM_CHECK_ERROR(options.mN >= options.mTileN, "N must be equal or larger than TileN."); + TLLM_CHECK_ERROR(options.mN % options.mTileN == 0, "N must be divisible by TileN."); + TLLM_CHECK_ERROR(!options.mTransposeMmaOutput, + "When batchM the MMA output has to be in row-major."); + } else { + TLLM_CHECK_ERROR(options.mM > 0 && options.mK > 0, "M and K must be larger than 0"); + TLLM_CHECK_ERROR(options.mM >= options.mTileM, "N must be equal or larger than tileN."); + TLLM_CHECK_ERROR(options.mM % options.mTileM == 0, "M must be divisible by TileM."); + TLLM_CHECK_ERROR(options.mTransposeMmaOutput, + "When batchN the MMA output has to be in column-major."); + } + } + + if (options.mUseDeepSeekFp8) { + if (batchM) { + // Make sure the GEMM-K dimension is a multiple of 128 when using DeepSeek FP8. + TLLM_CHECK_ERROR(options.mN % 128 == 0, + "GEMM-N must be a multiple of 128 when using DeepSeek Fp8. Found ", + options.mN); + } else { + // Make sure the GEMM-K dimension is a multiple of 128 when using DeepSeek FP8. + TLLM_CHECK_ERROR(options.mM % 128 == 0, + "GEMM-N must be a multiple of 128 when using DeepSeek Fp8. Found ", + options.mN); + } + // Make sure the GEMM-K dimension is a multiple of 128 when using DeepSeek FP8. + TLLM_CHECK_ERROR(options.mK % 128 == 0, + "GEMM-K must be a multiple of 128 when using DeepSeek Fp8. Found ", + options.mK); + + TLLM_CHECK_ERROR(options.mDtypeC != tg::Dtype::E2m1 && options.mDtypeA == tg::Dtype::E4m3 && + options.mDtypeB == tg::Dtype::E4m3, + "E2m1 is not supported with DeepSeek FP8"); + } + + if (batchM) { + if (options.mDtypeA == tg::Dtype::MxE2m1 && options.mMmaKind == tg::MmaKind::MxFp8Fp6Fp4) { + TLLM_CHECK_ERROR(doesRouteImplUseNoRoute(options.mRouteImpl), + "RouteAct is not supported with dtypeA = MxE2m1 and MxFp8Fp6Fp4."); + } + } else { + if (options.mDtypeB == tg::Dtype::MxE2m1 && options.mMmaKind == tg::MmaKind::MxFp8Fp6Fp4) { + TLLM_CHECK_ERROR(doesRouteImplUseNoRoute(options.mRouteImpl), + "RouteAct is not supported with dtypeB = MxE2m1 and MxFp8Fp6Fp4."); + } + } + + TLLM_CHECK_ERROR(options.mUseTmaStore, "Only TMA store is supported."); + if (batchM) { + TLLM_CHECK_ERROR(options.mLayoutA == gemm::MatrixLayout::MajorK, + "Activations must be in k-major format"); + } else { + TLLM_CHECK_ERROR(options.mLayoutB == gemm::MatrixLayout::MajorK, + "Activations must be in k-major format"); + } + + if (tg::mmaKindIsBlockFmt(options.mMmaKind) && !options.mUseDeepSeekFp8) { + if (!doesRouteImplUseNoRoute(options.mRouteImpl)) { + if (batchM) { + TLLM_CHECK_ERROR(options.mSfLayoutA == tg::SfLayout::Linear, + "Tokens need use SF linear layout when being routed"); + } else { + // Note: if B is cast from a non-block format to a block format, there are no SFs to load. + TLLM_CHECK_ERROR( + options.mSfLayoutB == tg::SfLayout::Linear || !tg::dtypeIsBlockFmt(options.mDtypeB), + "Tokens need use SF linear layout when being routed"); + } + } + + if (doesRouteImplUseTma(options.mRouteImpl)) { + TLLM_CHECK_ERROR(!batchM, "UTMALDG.GATHER4 only supported for batch N."); + + if (tg::mmaKindIsBlockFmt(options.mMmaKind)) { + auto dtypeRoute = batchM ? options.mDtypeA : options.mDtypeB; + TLLM_CHECK_ERROR(options.mTileK % tg::dtypeNumEltsPerSf(dtypeRoute) == 0, + "tileK needs to be a multiple of 16 * tg::dtypeNumEltsPerSf(dtypeA)."); + TLLM_CHECK_ERROR(options.mTileK % (tg::dtypeNumEltsPerSf(dtypeRoute) * 16) == 0, + "tileK needs to be a multiple of 16 * tg::dtypeNumEltsPerSf(dtypeA)."); + } + } + + if (!batchM || doesRouteImplUseNoRoute(options.mRouteImpl)) { + TLLM_CHECK_ERROR(options.mSfLayoutA == tg::SfLayout::R128c4, + "options.mSfLayoutA has to be tg::SfLayout::R128c4 when not being routed"); + } + } + + if (!gemm::isBiasTypeNone(options.mBiasType)) { + TLLM_CHECK_ERROR((gemm::isBiasTypeN(options.mBiasType) && + options.mBatchMode == BatchedGemmOptions::BatchMode::BatchM) || + (gemm::isBiasTypeM(options.mBiasType) && + options.mBatchMode == BatchedGemmOptions::BatchMode::BatchN), + "BatchedGemm supports only per channel bias."); + } + + // We do not handle the case where K is not a multiple of TileK. + // TMA based load handles the case transparently. + if (doesRouteImplUseLdgsts(options.mRouteImpl)) { + TLLM_CHECK_ERROR(options.mK % options.mTileK == 0, "K must be a multiple of TileK"); + } + + return isValid; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// +// +// BatchedGemmConfig +// +//////////////////////////////////////////////////////////////////////////////////////////////////// + +struct BatchedGemmConfig { + // When TRT-LLM Gen is exported to the other frameworks, the TLLM_GEN_EXPORT_INTERFACE must be + // defined. In this case, the cubins will be loaded from the provided data and function name. + // Otherwise, the kernel will be loaded from the CudaRunner. +#ifdef TLLM_GEN_EXPORT_INTERFACE + uint8_t const* mData{nullptr}; + uint32_t const mSize{0}; + uint32_t const mSharedMemSize{0}; + char const* mFunctionName{nullptr}; + uint32_t const mNumThreadsPerCTA{0}; + char const* mHash{nullptr}; +#else + trtllm::gen::CudaRunner* mCudaRunner{nullptr}; +#endif + + BatchedGemmOptions mOptions; + gemm::SmVersion mSm{gemm::SmVersion::Sm100a}; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline std::string dumpOptions(BatchedGemmOptions const& options) { + std::stringstream ss; + ss << gemmGatedAct::dumpOptions(options) << ", "; + ss << "mBatchedM={}," << std::endl; + ss << "mBatchedN={}," << std::endl; + ss << "mBatchMode=batchedGemm::BatchedGemmOptions::BatchMode(" + << static_cast(options.mBatchMode) << ")," << std::endl; + ss << "mNumBatches=" << options.mNumBatches << "," << std::endl; + ss << "mIsStaticBatch=" << options.mIsStaticBatch << "," << std::endl; + ss << "mNumTokens=" << options.mNumTokens << "," << std::endl; + ss << "mRouteImpl=batchedGemm::RouteImpl(" << static_cast(options.mRouteImpl) << ")," + << std::endl; + ss << "mGridWaitForPrimaryRouting=" << options.mGridWaitForPrimaryRouting << "," << std::endl; + ss << "mFusedAct=" << options.mFusedAct << "," << std::endl; + ss << "mNumRegsPerThreadNonEpilogueWarp=" << options.mNumRegsPerThreadNonEpilogueWarp << "," + << std::endl; + ss << "mNumRegsPerThreadEpilogueWarp=" << options.mNumRegsPerThreadEpilogueWarp << "," + << std::endl; + ss << "mNumRegsCastAWarps=" << options.mNumRegsCastAWarps << "," << std::endl; + ss << "mUseTmaOobOpt=" << options.mUseTmaOobOpt << std::endl; + return ss.str(); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace batchedGemm + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +#ifdef TLLM_GEN_EXPORT_INTERFACE + +#undef TLLM_CHECK_ERROR +#undef TLLM_CHECK_ERROR_FMT +#undef TLLM_CHECK_WARNING +#undef TLLM_LOG_WARNING +#undef TLLM_LOG_INFO +#undef TLLM_LOG_ERROR + +#endif // TLLM_GEN_EXPORT_INTERFACE + +} // namespace batchedGemm diff --git a/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/Enums.h b/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/Enums.h new file mode 100644 index 0000000000000000000000000000000000000000..6f2b1c270d5fae5e5eea690bf2ecc34126666b77 --- /dev/null +++ b/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/Enums.h @@ -0,0 +1,128 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & + * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include + +namespace batchedGemm { + +namespace gemm { + +enum class AllReduceAlgo : uint32_t { + // Does not apply all-reduce. + None = 0, + // Reduction occurs at L2 cache; pulls N-1 partial outputs from peer devices. Result is + // non-deterministic. Potentially lower latency at cost of higher memory traffic. + OneShot, + // Reduction occurs at switch; pulls 1/Nth of the output from switch (reduce-scatter phase) and + // store to multicast mem (all-gather phase). Result is deterministic. Lower memory traffic at + // cost of potentially higher latency. + TwoShot, +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +enum class MatrixLayout { + // K-major layout (default). [Mn, K] + MajorK = 0, + // M-major for A and N-major for B. [K, Mn] + MajorMn, + // Layout is blocked along the K dimension as seen in the diagram below. [K / blockK, Mn, blockK] + // where blockK is fixed at 128B + // + // ├────────────── K ──────────────┤ + // ┬ ┬ ├──── K block ───┤ + // │ │ │ 0 1 2 3 ║ 32 33 34 35 │ + // │ CTA0 │ 4 5 6 7 ║ 36 37 38 39 │ + // │ │ │ 8 9 10 11 ║ 40 41 42 43 │ + // │ ┴ │ 12 13 14 15 ║ 44 45 46 47 │ + // M ┬ ├────────────────║────────────────┤ + // │ │ │ 16 17 18 19 ║ 48 49 50 51 │ + // │ CTA1 │ 20 21 22 23 ║ 52 53 54 55 │ + // │ │ │ 24 25 26 27 ║ 56 57 58 59 │ + // ┴ ┴ │ 28 29 30 31 ║ 60 61 62 63 │ + BlockMajorK +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +enum class SplitK : uint32_t { + // No split-k is needed. I.e. mNumSlicesForSplitK == 1. + None = 0, + // CTAs computing one MN tile save partial results to global memory. + // Then wait on the barrier and the last CTA in the group loads partial results from gmem, + // sums them up and writes back to gmem. + Gmem, + // All CTAs in one CGA calculate partial sums. Then send the results to the smem of + // the last CTA in the CGA, which sums them up and writes to gmem. + Dsmem, +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +enum class BiasType : uint32_t { + // No bias. + None = 0, + // One bias value per N of the output tensor. + M = 1, + // One bias value per row M of the output tensor. + N = 2, + // One bias value for each element of the output tensor. + Mn = 3, +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +enum class TileScheduler { + // Static scheduler (Non-persistent). + Static = 0, + // Dynamic persistent scheduler. This is either based on an atomically incremented global work id + // prior to SM100 archs, or the HW supported work id scheduler based on UGETNEXTWORKID for SM100+. + Persistent, +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// Helper functions to check the SplitK type. + +#define SPLIT_K_FUNCTION(Mode) \ + inline bool doesSplitKUse##Mode(SplitK mode) { return (mode == SplitK::Mode); } + +SPLIT_K_FUNCTION(Gmem) +SPLIT_K_FUNCTION(Dsmem) + +#undef SPLIT_K_FUNCTION + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// Helper functions to check the Bias type. + +#define BIAS_TYPE_FUNCTION(Mode) \ + inline bool isBiasType##Mode(BiasType type) { return (type == BiasType::Mode); } + +BIAS_TYPE_FUNCTION(None) +BIAS_TYPE_FUNCTION(N) +BIAS_TYPE_FUNCTION(M) +BIAS_TYPE_FUNCTION(Mn) + +#undef BIAS_TYPE_FUNCTION + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace gemm + +} // namespace batchedGemm diff --git a/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/GemmGatedActOptions.h b/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/GemmGatedActOptions.h new file mode 100644 index 0000000000000000000000000000000000000000..d7cd564ae969cb4802a01dfceff196678629de10 --- /dev/null +++ b/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/GemmGatedActOptions.h @@ -0,0 +1,206 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & + * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include "GemmOptions.h" + +#ifdef TLLM_GEN_EXPORT_INTERFACE +#include + +#define TLLM_CHECK_ERROR(cond, ...) \ + if (!(cond)) { \ + printArgs(__VA_ARGS__); \ + return false; \ + } + +#define TLLM_LOG_ERROR(...) TLLM_CHECK_ERROR(false, __VA_ARGS__) + +#define TLLM_CHECK_ERROR_FMT(...) TLLM_CHECK_ERROR(false, __VA_ARGS__) + +#define TLLM_CHECK_WARNING(cond, ...) \ + if (!(cond)) { \ + printArgs(__VA_ARGS__); \ + return false; \ + } + +#define TLLM_LOG_WARNING(...) TLLM_CHECK_WARNING(false, __VA_ARGS__) + +#define TLLM_LOG_INFO(...) TLLM_CHECK_WARNING(false, __VA_ARGS__) + +#endif + +namespace batchedGemm { + +namespace gemmGatedAct { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace tg = trtllm::gen; + +// Type of the gated activation +enum class ActType { + // For ActType == SwiGlu, ideally we would like to have something like + // gatedAct = scaleC * (x0 * scaleAb + beta) * ((x1 * scaleGate) * sigmoid(alpha * x1 * + // scaleGate)). + // But for now, we use the simplified version + // gatedAct = scaleC' * (x0 + beta') * ((x1 * scaleGate) * sigmoid(alpha * x1 * scaleGate)), + // where x0 and x1 are the raw numbers from Gemm, while scaleC and scaleGate are input scales, + // beta' = beta / scaleAb, scaleC' = scaleC * scaleAb. + // + // GatedSilu is a special case of SwiGlu where the alpha is 1.0 and the beta is 0.0. + SwiGlu +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// Helper functions to check the ActType type. + +#define TLLM_ACT_TYPE_FUNCTION(actType) \ + inline bool is##actType(ActType type) { return (type == ActType::actType); } + +TLLM_ACT_TYPE_FUNCTION(SwiGlu) + +#undef TLLM_ACT_TYPE_FUNCTION + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline std::string getActTypeName(ActType type) { + switch (type) { + case ActType::SwiGlu: + return "SwiGlu"; + default: + return "Unknown type"; + } +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +struct GemmGatedActOptions : public gemm::GemmOptions { + GemmGatedActOptions() = default; + GemmGatedActOptions(gemm::GemmOptions options, ActType actType, bool clampBeforeAct) + : gemm::GemmOptions(options), mActType(actType), mClampBeforeAct(clampBeforeAct) {} + + // Type of the gated activation. + ActType mActType{ActType::SwiGlu}; + // Clamp the dequantized values to the range [-limit, limit]. + bool mClampBeforeAct{false}; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// Check if the options are valid or not. +inline bool checkAndUpdateGemmGatedActOptions(gemmGatedAct::GemmGatedActOptions& options, + bool isBlackwell, bool updateOptions = true) { + // tmpOut is already transposed at this stage + auto const hiddenSizeStr = options.mTransposeMmaOutput ? "M" : "N"; + auto const hiddenSize = options.mTransposeMmaOutput ? options.mM : options.mN; + auto const hiddenEpilogueTileSize = + options.mTransposeMmaOutput ? options.mEpilogueTileM : options.mEpilogueTileN; + + TLLM_CHECK_ERROR(hiddenSize % 2 == 0, hiddenSizeStr, " must be a multiple of 2."); + + TLLM_CHECK_ERROR((options.mTransposeMmaOutput ^ options.mUseShuffledMatrixA) == 0, + "Transpose mma output can only be used with shuffled A matrix. And vice versa."); + + if (options.mUseTmaStore) { + TLLM_CHECK_ERROR( + hiddenEpilogueTileSize * tg::dtypeGetNumBits(options.mDtypeC) / /* bits */ 8 % 32 == 0, + "Unsupported output hidden tile size"); + } + + if (options.mUseDeepSeekFp8) { + TLLM_CHECK_ERROR(hiddenSize % 256 == 0, "Output hidden size must be a multiple of 256"); + } + + if (options.mDtypeC == tg::Dtype::E2m1 || options.mDtypeC == tg::Dtype::MxE4m3) { + int const outHiddenSize = (options.mTransposeMmaOutput ? options.mM : options.mN) / 2; + int const hiddenGranularity = 4 * tg::dtypeNumEltsPerSf(options.mDtypeC); + TLLM_CHECK_ERROR(outHiddenSize % hiddenGranularity == 0, "Output hidden size (", outHiddenSize, + ") must be a multiple of ", hiddenGranularity, " for block-scaled outputs."); + } + + auto isValid = gemm::checkAndUpdateGemmOptions(options, isBlackwell, + /* tpGrpSize */ 1, updateOptions); + + if (!isValid) { + return false; + } + + if (options.mNumSlicesForSplitK > 1) { + TLLM_CHECK_ERROR(doesSplitKUseDsmem(options.mSplitK), + "Split-k GMEM and GemmGatedAct are not supported yet."); + } + + if (gemm::isBiasTypeMn(options.mBiasType)) { + TLLM_CHECK_ERROR(options.mTransposeMmaOutput, + "Bias type Mn is not supported with not transpose mma output."); + } + + return true; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline std::string dumpOptions(GemmGatedActOptions const& options) { + std::stringstream ss; + ss << gemm::dumpOptions(options) << ", "; + ss << "mActType=" + << "gemmGatedAct::ActType(" << static_cast(options.mActType) << ")," << std::endl; + ss << "mClampBeforeAct=" << options.mClampBeforeAct << "" << std::endl; + return ss.str(); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// +// +// GemmGatedActConfig +// +//////////////////////////////////////////////////////////////////////////////////////////////////// + +struct GemmGatedActConfig { + // When TRT-LLM Gen is exported to the other frameworks, the TLLM_GEN_EXPORT_INTERFACE must be + // defined. In this case, the cubins will be loaded from the provided data and function name. + // Otherwise, the kernel will be loaded from the CudaRunner. +#ifdef TLLM_GEN_EXPORT_INTERFACE + uint8_t const* mData{nullptr}; + uint32_t const mSize{0}; + uint32_t const mSharedMemSize{0}; + char const* mFunctionName{nullptr}; + uint32_t const mNumThreadsPerCTA{0}; + char const* mHash{nullptr}; +#else + trtllm::gen::CudaRunner* mCudaRunner{nullptr}; +#endif + + GemmGatedActOptions mOptions{}; + gemm::SmVersion mSm{gemm::SmVersion::Sm100a}; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace gemmGatedAct + +#ifdef TLLM_GEN_EXPORT_INTERFACE + +#undef TLLM_CHECK_ERROR +#undef TLLM_CHECK_ERROR_FMT +#undef TLLM_CHECK_WARNING +#undef TLLM_LOG_WARNING +#undef TLLM_LOG_INFO +#undef TLLM_LOG_ERROR +#endif // TLLM_GEN_EXPORT_INTERFACE + +} // namespace batchedGemm diff --git a/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/GemmOptions.h b/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/GemmOptions.h new file mode 100644 index 0000000000000000000000000000000000000000..1406ce2294ea6cf95d2f7a5c96adc16a8bb3126e --- /dev/null +++ b/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/GemmOptions.h @@ -0,0 +1,1217 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & + * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include +#include + +#include "Enums.h" +#include "KernelParams.h" +#include "KernelTraits.h" +#include "trtllm/gen/DtypeDecl.h" +#include "trtllm/gen/MmaDecl.h" +#include "trtllm/gen/SfLayoutDecl.h" +#ifndef TLLM_GEN_EXPORT_INTERFACE +#include "trtllm/gen/CudaRunner.h" +#include "trtllm/gen/GenCtx.h" +#else +#include + +template +void printArgs(T first, Args... args) { +#ifdef TLLM_GEN_DEBUG + std::cout << first; + if constexpr (sizeof...(args) > 0) { + std::cout << " "; + printArgs(args...); + } +#endif +} + +#define TLLM_CHECK_ERROR(cond, ...) \ + if (!(cond)) { \ + printArgs(__VA_ARGS__); \ + return false; \ + } + +#define TLLM_LOG_ERROR(...) TLLM_CHECK_ERROR(false, __VA_ARGS__) + +#define TLLM_CHECK_ERROR_FMT(cond, ...) TLLM_CHECK_ERROR(cond, __VA_ARGS__) + +#define TLLM_CHECK_WARNING(cond, ...) \ + if (!(cond)) { \ + printArgs(__VA_ARGS__); \ + return false; \ + } + +#define TLLM_LOG_WARNING(...) TLLM_CHECK_WARNING(false, __VA_ARGS__) + +#define TLLM_LOG_INFO(...) TLLM_CHECK_WARNING(false, __VA_ARGS__) + +#endif + +namespace batchedGemm { + +namespace gemm { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace tg = trtllm::gen; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// NOTE: when adding new parameters, please update the dumpOptions function and +// gemm_export_config.json for cubin export. +struct GemmOptions { +#ifndef TLLM_GEN_EXPORT_INTERFACE + // allow safely down-casting + virtual ~GemmOptions() = default; +#endif + + GemmOptions() = default; + GemmOptions(AllReduceAlgo allReduceAlgo, BiasType biasType, int blockK, int clusterDimX, + int clusterDimY, int clusterDimZ, tg::Dtype dtypeAcc, tg::Dtype dtypeA, + tg::Dtype dtypeB, tg::Dtype dtypeC, tg::Dtype dtypeMmaA, tg::Dtype dtypeMmaB, + bool enablesEarlyExit, bool enablesDelayedEarlyExit, bool enablesGlobalPtxKnobs, + int epilogueLdtmDps, int epilogueLdtmBits, int epilogueTileM, int epilogueTileN, + bool gridTriggerSecondaryA, bool gridTriggerSecondaryB, + bool gridWaitForPrimaryEarlyExit, bool gridWaitForPrimaryA, bool gridWaitForPrimaryB, + bool hoistLoadTaskInit, bool hoistMmaTaskTryWaits, int k, KernelTraits kernelTraits, + MatrixLayout layoutA, MatrixLayout layoutB, int m, int mmaK, tg::MmaKind mmaKind, + int mmaM, int mmaN, bool mockAllReduce, int n, int numSlicesForSplitK, + int numSlicesForSliceK, int numStages, int numStagesMma, + int numStagesMmaWithinWorkTile, int numStagesMmaAcrossWorkTile, int numStagesWorkId, + bool outputDebugTensors, bool patchF2fp, bool useShuffledMatrixA, bool sliceK, + SplitK splitK, bool transposeMmaOutput, int tileM, int tileN, int tileK, + bool useUnrollLoop2xForMma, bool useCustomMmaSchedule, + bool useHoistTryWaitForCustomMmaSchedule, bool useDeepSeekFp8, bool usePerTokenSfA, + bool usePerTokenSfB, bool useTmaStore, bool useTwoTmaLoadWarps, bool useTwoMmaWarps, + tg::SfLayout sfLayoutA, tg::SfLayout sfLayoutB, tg::SfLayout sfLayoutC, + int sfReshapeFactor, TileScheduler tileScheduler) + : mAllReduceAlgo{allReduceAlgo}, + mBiasType{biasType}, + mBlockK(blockK), + mClusterDimX{clusterDimX}, + mClusterDimY{clusterDimY}, + mClusterDimZ{clusterDimZ}, + mDtypeAcc{dtypeAcc}, + mDtypeA{dtypeA}, + mDtypeB{dtypeB}, + mDtypeC{dtypeC}, + mDtypeMmaA{dtypeMmaA}, + mDtypeMmaB{dtypeMmaB}, + mEnablesEarlyExit{enablesEarlyExit}, + mEnablesDelayedEarlyExit{enablesDelayedEarlyExit}, + mEnablesGlobalPtxKnobs{enablesGlobalPtxKnobs}, + mEpilogueLdtmDps{epilogueLdtmDps}, + mEpilogueLdtmBits{epilogueLdtmBits}, + mEpilogueTileM{epilogueTileM}, + mEpilogueTileN{epilogueTileN}, + mGridTriggerSecondaryA{gridTriggerSecondaryA}, + mGridTriggerSecondaryB{gridTriggerSecondaryB}, + mGridWaitForPrimaryEarlyExit{gridWaitForPrimaryEarlyExit}, + mGridWaitForPrimaryA{gridWaitForPrimaryA}, + mGridWaitForPrimaryB{gridWaitForPrimaryB}, + mHoistLoadTaskInit{hoistLoadTaskInit}, + mHoistMmaTaskTryWaits{hoistMmaTaskTryWaits}, + mK{k}, + mKernelTraits{kernelTraits}, + mLayoutA{layoutA}, + mLayoutB{layoutB}, + mM{m}, + mMmaK{mmaK}, + mMmaKind{mmaKind}, + mMmaM{mmaM}, + mMmaN{mmaN}, + mMockAllReduce{mockAllReduce}, + mN{n}, + mNumSlicesForSplitK{numSlicesForSplitK}, + mNumSlicesForSliceK{numSlicesForSliceK}, + mNumStages{numStages}, + mNumStagesMma{numStagesMma}, + mNumStagesMmaWithinWorkTile{numStagesMmaWithinWorkTile}, + mNumStagesMmaAcrossWorkTile{numStagesMmaAcrossWorkTile}, + mNumStagesWorkId{numStagesWorkId}, + mOutputDebugTensors{outputDebugTensors}, + mPatchF2fp{patchF2fp}, + mUseShuffledMatrixA{useShuffledMatrixA}, + mSliceK{sliceK}, + mSplitK{splitK}, + mTransposeMmaOutput{transposeMmaOutput}, + mTileM{tileM}, + mTileN{tileN}, + mTileK{tileK}, + mUseUnrollLoop2xForMma{useUnrollLoop2xForMma}, + mUseCustomMmaSchedule{useCustomMmaSchedule}, + mUseHoistTryWaitForCustomMmaSchedule{useHoistTryWaitForCustomMmaSchedule}, + mUseDeepSeekFp8{useDeepSeekFp8}, + mUsePerTokenSfA{usePerTokenSfA}, + mUsePerTokenSfB{usePerTokenSfB}, + mUseTmaStore{useTmaStore}, + mUseTwoTmaLoadWarps{useTwoTmaLoadWarps}, + mUseTwoMmaWarps{useTwoMmaWarps}, + mSfLayoutA{sfLayoutA}, + mSfLayoutB{sfLayoutB}, + mSfLayoutC{sfLayoutC}, + mSfReshapeFactor{sfReshapeFactor}, + mTileScheduler{tileScheduler} {} + + // The all-reduce algorithm. + AllReduceAlgo mAllReduceAlgo{AllReduceAlgo::None}; + // The type of bias. + BiasType mBiasType{BiasType::None}; + // Block size in the K dimension + int mBlockK{-1}; + // Cluster size in X dim. + int mClusterDimX{1}; + // Cluster size in Y dim. + int mClusterDimY{1}; + // Cluster size in Z dim. + int mClusterDimZ{1}; + // Data type of the accumulators. + tg::Dtype mDtypeAcc{tg::Dtype::Fp32}; + // Data type of the A matrix. + tg::Dtype mDtypeA{tg::Dtype::Fp16}; + // Data type of the B matrix. + tg::Dtype mDtypeB{tg::Dtype::Void}; + // Data type of the outputs. + tg::Dtype mDtypeC{tg::Dtype::Void}; + // Data type of the A matrix for the MMA, if different from the input type. + tg::Dtype mDtypeMmaA{tg::Dtype::Void}; + // Data type of the B matrix for the MMA, if different from the input type. + tg::Dtype mDtypeMmaB{tg::Dtype::Void}; + // Whether to enable early exit. + bool mEnablesEarlyExit{false}; + // Whether to enable delayed early exit to overlap + // numNonExitingCtas loading with the other instructions. + bool mEnablesDelayedEarlyExit{false}; + // Whether to enable the global PTX knobs for guiding the compiler optimizations. + bool mEnablesGlobalPtxKnobs{true}; + // The epilogue supports multiple LDTM shapes, although not every shape is applicable in every + // case. In particular: + // - On Hopper: must be 16dp256bit. + // - Transposed output: must be 16dp256bit. + // - Non-transposed output: + // - NvFp4 with fused activation: must be 32dp32bit. + // - Else it can be either 16dp256bit or 32dp32bit. + // The number of DP lanes in the epilogue LDTM. + int mEpilogueLdtmDps{16}; + // The number of bits in the epilogue LDTM. + int mEpilogueLdtmBits{256}; + // Tile size for the epilogue in M dimension. + int mEpilogueTileM{128}; + // Tile size for the epilogue in N dimension. + int mEpilogueTileN{32}; + // Whether load task A triggers the next grid. + bool mGridTriggerSecondaryA{false}; + // Whether load task B triggers the next grid. + bool mGridTriggerSecondaryB{false}; + // Whether the loads that check for an early exit should wait on a grid dependency. + bool mGridWaitForPrimaryEarlyExit{true}; + // Whether the load of A should wait on a grid dependency. + bool mGridWaitForPrimaryA{true}; + // Whether the load of B should wait on a grid dependency. + bool mGridWaitForPrimaryB{true}; + // Whether to hoist the initialization of the loading tasks. + bool mHoistLoadTaskInit{true}; + // Whether to hoist the mbarrier try_waits (e.g., mma.prodAcq, smemAb.consWait) in the MMA task. + bool mHoistMmaTaskTryWaits{false}; + // The K dimension of GEMM. + int mK{16 * 16}; + // Traits of the kernel. + KernelTraits mKernelTraits{}; + // Layout of A matrix + MatrixLayout mLayoutA{MatrixLayout::MajorK}; + // Layout of B matrix + MatrixLayout mLayoutB{MatrixLayout::MajorK}; + // The M dimension of GEMM. + int mM{128 * 2}; + // Size of the MMA instruction in the K dimension. + int mMmaK{16}; + // The kind of MMA instruction to use. + tg::MmaKind mMmaKind{tg::MmaKind::Auto}; + // Size of the MMA instruction in the M dimension. + int mMmaM{64}; + // Size of the MMA instruction in the N dimension. + int mMmaN{16}; + // Whether to mock all-reduce code for single-GPU debugging. + bool mMockAllReduce{false}; + // The N dimension of GEMM. + int mN{64 * 4}; + // Number of partitions along the K dimension. When mNumSlicesForSplitK > 1, + // the problem is distributed across several SMs, where each CTA works on its local K slice. + // Partial results are accumulated afterwards using either GMEM or DSMEM (in CGA) + // to exchange the data between CTAs. + int mNumSlicesForSplitK{1}; + // Number of slices for slice-K along K dimension. + int mNumSlicesForSliceK{1}; + // The depth of the mainloop pipeline. + int mNumStages{2}; + // The depth of the mma pipeline. Equals numStagesMmaWithinWorkTile * numStagesMmaAcrossWorkTile. + int mNumStagesMma{1}; + // The depth of the mma pipeline within work tile. Only GmemC classes with "WithAccInReg" suffix + // are allowed to be greater than 1. + int mNumStagesMmaWithinWorkTile{-1}; + // The depth of the mma pipeline across work tiles in the persistent loop. + int mNumStagesMmaAcrossWorkTile{-1}; + // The depth of the work id pipeline and the work throttle pipeline. + int mNumStagesWorkId{3}; + // Whether to output debug tensors. + bool mOutputDebugTensors{false}; + // Patch float conversions. + bool mPatchF2fp{false}; + // Reorder rows/cols in the A matrix for the better memory accesses in the M-major epilogue. + bool mUseShuffledMatrixA{false}; + // Slice-K implementation to use TileM dimension for TileK. + bool mSliceK{false}; + // The location of the exchange for split-K (it's None when split-K is disabled). + SplitK mSplitK{SplitK::None}; + // Save output of MMA in M-major format. + bool mTransposeMmaOutput{false}; + // M tile dimension of GEMM. + int mTileM{128}; + // N tile dimension of GEMM. + int mTileN{32}; + // K tile dimension of GEMM. + int mTileK{16}; + // Whether to unroll the loop by 2x. + bool mUseUnrollLoop2xForMma{true}; + // Use custom MMA schedule optimized for low-latency. + bool mUseCustomMmaSchedule{false}; + // The purpose of hoisting trywaits is to opportunistically peek at the availability of the next + // k-block. It benefits when the next k-block is already available and thus sustaining the + // momentum, but it adds latency to the first k-block for smaller k-loop. + bool mUseHoistTryWaitForCustomMmaSchedule{false}; + // Use DeepSeek Fp8. + bool mUseDeepSeekFp8{false}; + // Apply per-token scales from A + bool mUsePerTokenSfA{false}; + // Apply per-token scales from B + bool mUsePerTokenSfB{false}; + // Use TMA to store the result. + bool mUseTmaStore{true}; + // Use two different warps for A and B matrix load. + bool mUseTwoTmaLoadWarps{false}; + // Use two different warps for MMA tasks. Applicable only to DeepSeek FP8. + bool mUseTwoMmaWarps{false}; + // Scale factors layout for A. + tg::SfLayout mSfLayoutA{tg::SfLayout::R128c4}; + // Scale factors layout for B. + tg::SfLayout mSfLayoutB{tg::SfLayout::R128c4}; + // Scale factors layout for C. + tg::SfLayout mSfLayoutC{tg::SfLayout::R128c4}; + // Number of "repeats", i.e. reshaping factor, to fold hidden dimension into SfBlock dimension. + // As result, the hidden dimension of the SF tensor must be a multiple of NumRepeats * + // numEltsPerSf * 4. This reduces the problem shape space that the kernel is able to run. + // But it reduces the number of L2 requests under the hood and potentially improves perf. + // Applies to layout 8x4 only. + int mSfReshapeFactor{1}; + // Tile scheduler type. + TileScheduler mTileScheduler{TileScheduler::Static}; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +enum class SmVersion { Sm90a, Sm100a }; + +//////////////////////////////////////////////////////////////////////////////////////////////////// +// +// GemmConfig +// +//////////////////////////////////////////////////////////////////////////////////////////////////// + +struct GemmConfig { + // When TRT-LLM Gen is exported to the other frameworks, the TLLM_GEN_EXPORT_INTERFACE must be + // defined. In this case, the cubins will be loaded from the provided data and function name. + // Otherwise, the kernel will be loaded from the CudaRunner. +#ifdef TLLM_GEN_EXPORT_INTERFACE + uint8_t const* mData{nullptr}; + uint32_t const mSize{0}; + uint32_t const mSharedMemSize{0}; + char const* mFunctionName{nullptr}; + uint32_t const mNumThreadsPerCTA{0}; + char const* mHash{nullptr}; +#else + trtllm::gen::CudaRunner* mCudaRunner{nullptr}; +#endif + + GemmOptions mOptions{}; + SmVersion mSm{SmVersion::Sm100a}; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// Serialization helpers. +template +inline std::string toString(T e) { + return std::to_string(e); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +inline std::string toString(trtllm::gen::Dtype e) { + return trtllm::gen::dtypeToString(e); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +inline std::string toString(trtllm::gen::MmaKind e) { + return trtllm::gen::mmaKindToString(e); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline std::string dumpOptions(GemmOptions const& options) { + std::stringstream ss; + ss << "mAllReduceAlgo=" + << "gemm::AllReduceAlgo(" << static_cast(options.mAllReduceAlgo) << ")" + << "," << std::endl; + ss << "mBiasType=" + << "gemm::BiasType(" << static_cast(options.mBiasType) << ")" + << "," << std::endl; + ss << "mBlockK=" << options.mBlockK << "," << std::endl; + ss << "mClusterDimX=" << options.mClusterDimX << "," << std::endl; + ss << "mClusterDimY=" << options.mClusterDimY << "," << std::endl; + ss << "mClusterDimZ=" << options.mClusterDimZ << "," << std::endl; + ss << "mDtypeAcc=" + << "trtllm::gen::Dtype(" << static_cast(options.mDtypeAcc) << ")" + << "," << std::endl; + ss << "mDtypeA=" + << "trtllm::gen::Dtype(" << static_cast(options.mDtypeA) << ")" + << "," << std::endl; + ss << "mDtypeB=" + << "trtllm::gen::Dtype(" << static_cast(options.mDtypeB) << ")" + << "," << std::endl; + ss << "mDtypeC=" + << "trtllm::gen::Dtype(" << static_cast(options.mDtypeC) << ")" + << "," << std::endl; + ss << "mDtypeMmaA=" + << "trtllm::gen::Dtype(" << static_cast(options.mDtypeMmaA) << ")" + << "," << std::endl; + ss << "mDtypeMmaB=" + << "trtllm::gen::Dtype(" << static_cast(options.mDtypeMmaB) << ")" + << "," << std::endl; + ss << "mEnablesEarlyExit=" << options.mEnablesEarlyExit << "," << std::endl; + ss << "mEnablesDelayedEarlyExit=" << options.mEnablesDelayedEarlyExit << "," << std::endl; + ss << "mEnablesGlobalPtxKnobs=" << options.mEnablesGlobalPtxKnobs << "," << std::endl; + ss << "mEpilogueLdtmDps=" << options.mEpilogueLdtmDps << "," << std::endl; + ss << "mEpilogueLdtmBits=" << options.mEpilogueLdtmBits << "," << std::endl; + ss << "mEpilogueTileM=" << options.mEpilogueTileM << "," << std::endl; + ss << "mEpilogueTileN=" << options.mEpilogueTileN << "," << std::endl; + ss << "mGridTriggerSecondaryA=" << options.mGridTriggerSecondaryA << "," << std::endl; + ss << "mGridTriggerSecondaryB=" << options.mGridTriggerSecondaryB << "," << std::endl; + ss << "mGridWaitForPrimaryEarlyExit=" << options.mGridWaitForPrimaryEarlyExit << "," << std::endl; + ss << "mGridWaitForPrimaryA=" << options.mGridWaitForPrimaryA << "," << std::endl; + ss << "mGridWaitForPrimaryB=" << options.mGridWaitForPrimaryB << "," << std::endl; + ss << "mHoistLoadTaskInit=" << options.mHoistLoadTaskInit << "," << std::endl; + ss << "mHoistMmaTaskTryWaits=" << options.mHoistMmaTaskTryWaits << "," << std::endl; + ss << "mK=" << options.mK << "," << std::endl; + ss << "mKernelTraits={}" + << "," << std::endl; + ss << "mLayoutA=gemm::MatrixLayout(" << static_cast(options.mLayoutA) << ")" + << "," << std::endl; + ss << "mLayoutB=gemm::MatrixLayout(" << static_cast(options.mLayoutB) << ")" + << "," << std::endl; + ss << "mM=" << options.mM << "," << std::endl; + ss << "mMmaK=" << options.mMmaK << "," << std::endl; + ss << "mMmaKind=" + << "trtllm::gen::MmaKind(" << static_cast(options.mMmaKind) << ")" + << "," << std::endl; + ss << "mMmaM=" << options.mMmaM << "," << std::endl; + ss << "mMmaN=" << options.mMmaN << "," << std::endl; + ss << "mMockAllReduce=" << options.mMockAllReduce << "," << std::endl; + ss << "mN=" << options.mN << "," << std::endl; + ss << "mNumSlicesForSplitK=" << options.mNumSlicesForSplitK << "," << std::endl; + ss << "mNumSlicesForSliceK=" << options.mNumSlicesForSliceK << "," << std::endl; + ss << "mNumStages=" << options.mNumStages << "," << std::endl; + ss << "mNumStagesMma=" << options.mNumStagesMma << "," << std::endl; + ss << "mNumStagesMmaWithinWorkTile=" << options.mNumStagesMmaWithinWorkTile << "," << std::endl; + ss << "mNumStagesMmaAcrossWorkTile=" << options.mNumStagesMmaAcrossWorkTile << "," << std::endl; + ss << "mNumStagesWorkId=" << options.mNumStagesWorkId << "," << std::endl; + ss << "mOutputDebugTensors=" << options.mOutputDebugTensors << "," << std::endl; + ss << "mPatchF2fp=" << options.mPatchF2fp << "," << std::endl; + ss << "mUseShuffledMatrixA=" << options.mUseShuffledMatrixA << "," << std::endl; + ss << "mSliceK=" << options.mSliceK << "," << std::endl; + ss << "mSplitK=" + << "gemm::SplitK(" << static_cast(options.mSplitK) << ")" + << "," << std::endl; + ss << "mTransposeMmaOutput=" << options.mTransposeMmaOutput << "," << std::endl; + ss << "mTileM=" << options.mTileM << "," << std::endl; + ss << "mTileN=" << options.mTileN << "," << std::endl; + ss << "mTileK=" << options.mTileK << "," << std::endl; + ss << "mUseUnrollLoop2xForMma=" << options.mUseUnrollLoop2xForMma << "," << std::endl; + ss << "mUseCustomMmaSchedule=" << options.mUseCustomMmaSchedule << "," << std::endl; + ss << "mUseHoistTryWaitForCustomMmaSchedule=" << options.mUseHoistTryWaitForCustomMmaSchedule + << "," << std::endl; + ss << "mUseDeepSeekFp8=" << options.mUseDeepSeekFp8 << "," << std::endl; + ss << "mUsePerTokenSfA=" << options.mUsePerTokenSfA << "," << std::endl; + ss << "mUsePerTokenSfB=" << options.mUsePerTokenSfB << "," << std::endl; + ss << "mUseTmaStore=" << options.mUseTmaStore << "," << std::endl; + ss << "mUseTwoTmaLoadWarps=" << options.mUseTwoTmaLoadWarps << "," << std::endl; + ss << "mUseTwoMmaWarps=" << options.mUseTwoMmaWarps << "," << std::endl; + ss << "mSfLayoutA=" + << "trtllm::gen::SfLayout(" << static_cast(options.mSfLayoutA) << ")" + << "," << std::endl; + ss << "mSfLayoutB=" + << "trtllm::gen::SfLayout(" << static_cast(options.mSfLayoutB) << ")" + << "," << std::endl; + ss << "mSfLayoutC=" + << "trtllm::gen::SfLayout(" << static_cast(options.mSfLayoutC) << ")" + << "," << std::endl; + ss << "mSfReshapeFactor=" << options.mSfReshapeFactor << "," << std::endl; + ss << "mTileScheduler=" + << "gemm::TileScheduler(" << static_cast(options.mTileScheduler) << ")" << std::endl; + return ss.str(); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +inline T divUp(T a, T b) { + return (a + b - 1) / b; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +inline T divUpMul(T a, T b) { + return gemm::divUp(a, b) * b; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline int32_t getShuffleBlockSize(int epilogueTileM) { + int shuffleBlockSize = 16; + if (epilogueTileM % 128 == 0) { + shuffleBlockSize = 32; + } + return shuffleBlockSize; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// Check if the options are valid or not. +inline bool checkAndUpdateGemmOptions(GemmOptions& options, bool isBlackwell, int /* tpGrpSize */, + bool updateOptions = true) { + if (options.mDtypeB == tg::Dtype::Void) { + if (updateOptions) { + options.mDtypeB = options.mDtypeA; + } else { + return false; + } + } + + // If not specified, used the input dtypes as MMA dtypes (no cast required). + if (options.mDtypeMmaA == tg::Dtype::Void) { + if (updateOptions) { + options.mDtypeMmaA = options.mDtypeA; + } else { + return false; + } + } + if (options.mDtypeMmaB == tg::Dtype::Void) { + if (updateOptions) { + options.mDtypeMmaB = options.mDtypeB; + } else { + return false; + } + } + + // Check that the A cast is supported. + // Currently, we only support {MxFp4, NvFp4} -> Bf16. + TLLM_CHECK_ERROR( + (options.mDtypeA == options.mDtypeMmaA) || + ((options.mDtypeA == tg::Dtype::MxE2m1 || options.mDtypeA == tg::Dtype::E2m1) && + options.mDtypeMmaA == tg::Dtype::Bfloat16) || + (options.mDtypeA == tg::Dtype::E2m1 && options.mDtypeMmaA == tg::Dtype::E4m3), + "Unsupported cast for A: ", tg::dtypeToString(options.mDtypeA), " -> ", + tg::dtypeToString(options.mDtypeMmaA)); + + // Check that the B cast is supported. + // Currently, we only support Fp8 -> MxFp8. + // TODO: add same support for A (no transpose) + TLLM_CHECK_ERROR( + (options.mDtypeB == options.mDtypeMmaB) || + (options.mDtypeB == tg::Dtype::E4m3 && options.mDtypeMmaB == tg::Dtype::MxE4m3), + "Unsupported cast for B: ", tg::dtypeToString(options.mDtypeB), " -> ", + tg::dtypeToString(options.mDtypeMmaB)); + + if (options.mDtypeA != options.mDtypeMmaA) { + TLLM_CHECK_ERROR(options.mTileM == 128, + "TileM must be 128 when casting the input matrix A before the MMA."); + } + + if (options.mPatchF2fp) { + TLLM_CHECK_ERROR( + options.mDtypeA == tg::Dtype::MxE2m1 && options.mDtypeMmaA == tg::Dtype::Bfloat16, + "PatchF2fp is only supported for MxFp4 to Bf16 casts."); + } + + // FIXME: We do not support different dtypes for A and B when not on Blackwell. + if (!isBlackwell) { + TLLM_CHECK_ERROR(options.mDtypeMmaA == options.mDtypeMmaB, + "For non-Blackwell, A and B must have the same dtype."); + } + + // Check that the different dtypes for A and B are supported by the tensor core + // kind::f8f6f4 + if (options.mDtypeMmaA == tg::Dtype::E4m3 || options.mDtypeMmaA == tg::Dtype::E2m1) { + TLLM_CHECK_ERROR(options.mDtypeMmaB == tg::Dtype::E4m3 || options.mDtypeMmaB == tg::Dtype::E2m1, + "For dtypeMmaA = E4m3/E2m1 A, dtypeMmaB must also be E4m3/E2m1."); + } + + // kind::mxf8f6f4 + if (options.mDtypeMmaA == tg::Dtype::MxE4m3 || options.mDtypeMmaA == tg::Dtype::MxE2m1) { + TLLM_CHECK_ERROR( + options.mDtypeMmaB == tg::Dtype::MxE4m3 || options.mDtypeMmaB == tg::Dtype::MxE2m1, + "For dtypeMmaA = MxE4m3 or MxE2m1, dtypeMmaB must also be MxE4m3 or MxE2m1."); + } + if (options.mDtypeMmaB == tg::Dtype::MxE4m3 || options.mDtypeMmaB == tg::Dtype::MxE2m1) { + TLLM_CHECK_ERROR( + options.mDtypeMmaA == tg::Dtype::MxE4m3 || options.mDtypeMmaA == tg::Dtype::MxE2m1, + "For dtypeMmaB = MxE4m3 or MxE2m1, dtypeMmaA must also be MxE4m3 or MxE2m1."); + } + + // kind::f16 + if (options.mDtypeMmaA == tg::Dtype::Fp16 || options.mDtypeMmaA == tg::Dtype::Bfloat16) { + TLLM_CHECK_ERROR(options.mDtypeMmaB == options.mDtypeMmaA, + "For dtypeMmaA = Fp16/Bfloat16, dtypeMmaB must be the same as dtypeMmaA."); + } + + // When one of the inputs needs to be cast, we must use two load warps. + if ((options.mDtypeMmaA != options.mDtypeA || options.mDtypeMmaB != options.mDtypeB) && + !options.mUseTwoTmaLoadWarps) { + TLLM_LOG_WARNING("Two TMA load warps must be enabled if any of the inputs needs to be cast."); + } + + // When different dtypes are used for A and B, we must use different tiles to do the loading. + // It is not strictly required, but current implementation of SmemAb requires that. + if (options.mDtypeA != options.mDtypeB) { + TLLM_CHECK_ERROR(options.mUseTwoTmaLoadWarps, + "Two TMA load warps must be enabled for different input types of A and B."); + } + + // Get the mma kind for the input types. + if (options.mMmaKind == tg::MmaKind::Auto) { + if (updateOptions) { + options.mMmaKind = dtypeGetMmaKind(options.mDtypeMmaA, options.mDtypeMmaB); + } else { + return false; + } + } + + if ((options.mMmaKind == tg::MmaKind::Fp8Fp6Fp4 || + options.mMmaKind == tg::MmaKind::MxFp8Fp6Fp4) && + options.mMmaK != 32) { + TLLM_LOG_WARNING("Unsupported MmaK (", options.mMmaK, + ") for MmaKind=", gemm::toString(options.mMmaKind), ". Setting MmaK to 32"); + if (updateOptions) { + options.mMmaK = 32; + options.mTileK = std::max(options.mMmaK, options.mTileK); + } else { + return false; + } + } + + // Check LDTM shape. + if (isBlackwell) { + TLLM_CHECK_ERROR((options.mEpilogueLdtmDps == 16 && options.mEpilogueLdtmBits == 256) || + (options.mEpilogueLdtmDps == 32 && options.mEpilogueLdtmBits == 32), + "Unsupported LDTM shape: ", options.mEpilogueLdtmDps, "dp", + options.mEpilogueLdtmBits, "bit."); + if (options.mEpilogueTileM == 64) { + TLLM_CHECK_ERROR(options.mEpilogueLdtmDps == 16, + "Unsupported LDTM shape for epilogueTileM=64: ", options.mEpilogueLdtmDps, + "dp", options.mEpilogueLdtmBits, "bit."); + } + if (options.mTransposeMmaOutput) { + // We can't use 32dp32bit LDTM for transposed outputs because we need each thread to own + // multiple consecutive output elements. + TLLM_CHECK_ERROR((options.mEpilogueLdtmDps == 16 && options.mEpilogueLdtmBits == 256), + "Only 16dp256bit LDTM is supported for transposed outputs."); + } + } else { + TLLM_CHECK_ERROR( + options.mEpilogueLdtmDps == 16 && options.mEpilogueLdtmBits == 256, + "Hopper does not use TMEM. The register layout corresponds to 16dp256bit. Got ", + options.mEpilogueLdtmDps, "dp", options.mEpilogueLdtmBits, "bit."); + } + + // Constraints for NvFp4 and MxFp8. + if ((options.mMmaKind == tg::MmaKind::MxFp4NvFp4 || + options.mMmaKind == tg::MmaKind::MxFp8Fp6Fp4 || options.mDtypeC == tg::Dtype::MxE4m3) && + options.mMmaM != 128) { + // MMA M must be 128 when the input uses block scaling, or when the output is an Mx format. + int newTileM = 128 * divUp(options.mTileM, 128); + TLLM_LOG_WARNING("Unsupported MmaM (", options.mMmaM, + ") for MmaKind=", gemm::toString(options.mMmaKind), + ". Setting MmaM to 128 and TileM to ", newTileM); + if (updateOptions) { + options.mMmaM = 128; + options.mTileM = newTileM; + } else { + return false; + } + } + if (options.mMmaKind == tg::MmaKind::MxFp4NvFp4 || options.mMmaKind == tg::MmaKind::MxFp8Fp6Fp4) { + TLLM_CHECK_ERROR(isBlackwell, "Block scaling is only supported on Blackwell"); + + int const mmaK = (options.mMmaKind == tg::MmaKind::MxFp4NvFp4) ? 64 : 32; + if (options.mMmaK != mmaK) { + int newTileK = mmaK * divUp(options.mTileK, mmaK); + TLLM_LOG_WARNING("Unsupported MmaK (", options.mMmaK, + ") for MmaKind=", gemm::toString(options.mMmaKind), ". Setting MmaK to ", + mmaK, " and TileK to ", newTileK); + if (updateOptions) { + options.mMmaK = mmaK; + options.mTileK = newTileK; + } else { + return false; + } + } + + // The MMA N may only be smaller than 64 if it is equal to the tile N. + TLLM_CHECK_ERROR(options.mMmaN >= 64 || options.mMmaN == options.mTileN, "MmaN (", + options.mMmaN, ") must be >= 64 or equal to TileN (", options.mTileN, ")"); + } + if (tg::dtypeIsBlockFmt(options.mDtypeA)) { + int numEltsPerSfA = tg::dtypeNumEltsPerSf(options.mDtypeA); + TLLM_CHECK_ERROR(options.mTileK % (4 * numEltsPerSfA) == 0, "TileK (", options.mTileK, + ") must be a multiple of ", (4 * numEltsPerSfA), " for typeA ", + gemm::toString(options.mDtypeA)); + auto const numEltsPerSfAInK = options.mK / numEltsPerSfA; + TLLM_CHECK_ERROR(numEltsPerSfAInK % 4 == 0, "K dimension of scaling factors for A (", + numEltsPerSfAInK, ") must be a multiple of 4"); + } + if (tg::dtypeIsBlockFmt(options.mDtypeB)) { + TLLM_CHECK_ERROR(options.mSfLayoutB == tg::SfLayout::R128c4 || + options.mSfLayoutB == tg::SfLayout::R8c4 || + options.mSfLayoutB == tg::SfLayout::Linear, + "Only the 128x4 and 8x4 SF layouts are supported for B, got ", + tg::sfLayoutToString(options.mSfLayoutB)); + + // TileN must be a multiple of the number of rows per SF tile. + int const numSfTileRowsB = options.mSfLayoutB == tg::SfLayout::R128c4 ? 128 : 8; + TLLM_CHECK_ERROR(options.mTileN % numSfTileRowsB == 0, "TileN (", options.mTileN, + ") must be a multiple of ", numSfTileRowsB, " for B SF layout ", + tg::sfLayoutToString(options.mSfLayoutB)); + + int numEltsPerSfB = tg::dtypeNumEltsPerSf(options.mDtypeB); + TLLM_CHECK_ERROR(options.mTileK % (4 * numEltsPerSfB) == 0, "TileK (", options.mTileK, + ") must be a multiple of ", (4 * numEltsPerSfB), " for typeB ", + gemm::toString(options.mDtypeB)); + auto const numEltsPerSfBInK = options.mK / numEltsPerSfB; + TLLM_CHECK_ERROR(numEltsPerSfBInK % 4 == 0, "K dimension of scaling factors for B (", + numEltsPerSfBInK, ") must be a multiple of 4"); + } + + int32_t padMultiplierA = 1; + int32_t padMultiplierB = 1; + if (options.mMmaKind == tg::MmaKind::MxFp8Fp6Fp4) { + if (options.mDtypeA == tg::Dtype::MxE2m1) { + padMultiplierA = 2; + } + if (options.mDtypeB == tg::Dtype::MxE2m1) { + padMultiplierB = 2; + } + } + TLLM_CHECK_ERROR( + (padMultiplierA * tg::dtypeGetNumBits(options.mDtypeA) * options.mK / 8) % 16 == 0, + "K dimension of A must be aligned to 16 bytes."); + TLLM_CHECK_ERROR( + (padMultiplierB * tg::dtypeGetNumBits(options.mDtypeB) * options.mK / 8) % 16 == 0, + "K dimension of B must be aligned to 16 bytes."); + + if (options.mDtypeC == tg::Dtype::E2m1 || options.mDtypeC == tg::Dtype::MxE4m3) { + TLLM_CHECK_ERROR(isBlackwell, "Block scaling is only supported on Blackwell"); + + TLLM_CHECK_ERROR( + options.mSfLayoutC == tg::SfLayout::R128c4 || options.mSfLayoutC == tg::SfLayout::R8c4, + "Only the 128x4 and 8x4 SF layouts are supported for C."); + int const numSfTileRowsC = options.mSfLayoutC == tg::SfLayout::R128c4 ? 128 : 8; + int const tileTokenDim = options.mTransposeMmaOutput ? options.mTileN : options.mTileM; + TLLM_CHECK_ERROR_FMT(tileTokenDim % numSfTileRowsC == 0, + "Tile%s (%d) must be a multiple of %d for C SF layout %s", + options.mTransposeMmaOutput ? "N" : "M", tileTokenDim, numSfTileRowsC, + tg::sfLayoutToString(options.mSfLayoutC).c_str()); + + int const hiddenDim = options.mTransposeMmaOutput ? options.mM : options.mN; + int const hiddenGranularity = 4 * tg::dtypeNumEltsPerSf(options.mDtypeC); + TLLM_CHECK_ERROR(hiddenDim % hiddenGranularity == 0, "Hidden dim (", hiddenDim, + ") must be a multiple of ", hiddenGranularity, " for block-scaled outputs."); + TLLM_CHECK_ERROR(!options.mTransposeMmaOutput || options.mUseShuffledMatrixA, + "Transposing block-scaled outputs requires shuffled A."); + } + + // If dtypeC is unspecified (Dtype::Void), assign to the input dtype. + if (options.mDtypeC == tg::Dtype::Void) { + TLLM_LOG_INFO("Setting dtypeC to ", tg::dtypeToString(options.mDtypeA)); + if (updateOptions) { + options.mDtypeC = options.mDtypeA; + } else { + return false; + } + } + + // Set epilogue tile sizes to the output tile sizes, when epilogue tile sizes are incorrect. + if (options.mTileM % options.mEpilogueTileM != 0) { + TLLM_LOG_WARNING("TileM (", options.mTileM, ") must be divisible by EpilogueTileM (", + options.mEpilogueTileM, "). Setting EpilogueTileM to TileM"); + if (updateOptions) { + options.mEpilogueTileM = options.mTileM; + } else { + return false; + } + } + + if (options.mTileN % options.mEpilogueTileN != 0) { + TLLM_LOG_WARNING("TileN (", options.mTileN, ") must be divisible by EpilogueTileN (", + options.mEpilogueTileN, "). Setting EpilogueTileN to TileN"); + if (updateOptions) { + options.mEpilogueTileN = options.mTileN; + } else { + return false; + } + } + + // On Hopper, epilogue tile sizes are the same as output tiles. + if (!isBlackwell && + (options.mEpilogueTileM != options.mTileM || options.mEpilogueTileN != options.mTileN)) { + TLLM_LOG_WARNING( + "Overwriting epilogueTileM and epilogueTileN to match tileM and tileN respectively"); + if (updateOptions) { + options.mEpilogueTileM = options.mTileM; + options.mEpilogueTileN = options.mTileN; + } else { + return false; + } + } + + // Unsupported epilogue tile size. + if (options.mMmaM == 128 && options.mEpilogueTileM != options.mTileM) { + TLLM_LOG_WARNING( + "When MmaM = 128, EpilogueTileM must be equal to TileM. Setting EpilogueTileM to TileM"); + if (updateOptions) { + options.mEpilogueTileM = options.mTileM; + } else { + return false; + } + } + + TLLM_CHECK_ERROR(options.mM > 0 && options.mN > 0 && options.mK > 0, + "M, N and K must be larger than 0"); + TLLM_CHECK_ERROR(options.mNumSlicesForSplitK > 0, "Split K must be larger than 0."); + + if (options.mUseShuffledMatrixA) { + auto const shuffleBlockSize = getShuffleBlockSize(options.mEpilogueTileM); + TLLM_CHECK_ERROR(options.mM % shuffleBlockSize == 0, + "M must be a multiple of shuffle block size (", shuffleBlockSize, + ") when useShuffledMatrixA"); + } + + if (!options.mSliceK) { + TLLM_CHECK_ERROR(options.mMmaM <= options.mEpilogueTileM, + "EpilogueTileM must be larger or equal than mmaM."); + } + TLLM_CHECK_ERROR( + options.mTileM % options.mEpilogueTileM == 0 && options.mTileN % options.mEpilogueTileN == 0, + "TileM and TileN must be divisible by EpilogueTileM and EpilogueTileN respectively."); + TLLM_CHECK_ERROR(options.mClusterDimX == 1 && options.mClusterDimY == 1, + "GEMM does not support cluster in X and Y dimensions."); + TLLM_CHECK_ERROR(options.mClusterDimZ == 1 || options.mNumSlicesForSplitK > 1, + "Cluster DimZ is only allowed for split-k."); + TLLM_CHECK_ERROR(options.mTileM <= 128, "GEMM does not support TileM > 128."); + + // When the A-matrix is shuffled, the output must be transposed. + if (options.mUseShuffledMatrixA) { + // TODO add matrix shuffle for N-major epilogue. + TLLM_CHECK_ERROR( + options.mTransposeMmaOutput, + "Shuffled matrix A is only supported with M-major epilogue. Set -transposeMmaOutput"); + } + + // Check all-reduce options. + if (options.mAllReduceAlgo == AllReduceAlgo::OneShot) { + // One shot is implemented with PTX cp.reduce.async.bulk.tensor which supports only the + // following types for reduce add: u32, s32, u64, f32, f16, bf16. + // + // See: https://docs.nvidia.com/cuda/parallel-thread-execution/ + // #data-movement-and-conversion-instructions-cp-reduce-async-bulk-tensor + std::set dtypeSupported{tg::Dtype::UInt32, tg::Dtype::Int32, tg::Dtype::UInt64, + tg::Dtype::Fp32, tg::Dtype::Fp16, tg::Dtype::Bfloat16}; + TLLM_CHECK_ERROR(dtypeSupported.find(options.mDtypeC) != dtypeSupported.end(), + "Unsupported output dtype ", tg::dtypeToString(options.mDtypeC)); + } else if (options.mAllReduceAlgo == AllReduceAlgo::TwoShot) { + // TODO(anchengc): + // Input dtype == output dtype -> can perform all-reduce in-place. + // Input dtype != output dtype -> must perform all-reduce out of place. + TLLM_CHECK_ERROR_FMT( + options.mDtypeC == options.mDtypeAcc, + "Not implemented - mixed dtype (dtypeC (%s) != dtypeAcc (%s)) requires out of place update", + tg::dtypeToString(options.mDtypeC).c_str(), tg::dtypeToString(options.mDtypeAcc).c_str()); + } + if (options.mAllReduceAlgo != AllReduceAlgo::None) { + TLLM_CHECK_ERROR(options.mUseTmaStore, "Non-TMA store with all-reduce is not implemented"); + } + + if (updateOptions) { + if (options.mNumSlicesForSplitK == 1) { + // No split-k. + options.mSplitK = SplitK::None; + } else if (options.mNumSlicesForSplitK > 1 && options.mClusterDimZ == 1) { + // Split-k with exchange through gmem. + options.mSplitK = SplitK::Gmem; + } else { + // Split-k with exchange through Dsmem. + options.mSplitK = SplitK::Dsmem; + } + } + // For GMEM-based split-K, we write 4 elements at once. + if (options.mSplitK == SplitK::Gmem) { + TLLM_CHECK_ERROR((options.mM * options.mN) % 4 == 0, + "M * N must be a multiple of 4 for Split-K"); + } + + if (options.mNumSlicesForSplitK > 1) { + if ((options.mEpilogueTileM != options.mTileM || options.mEpilogueTileN != options.mTileN) && + !options.mUseDeepSeekFp8) { + TLLM_LOG_WARNING( + "Overwriting epilogueTileM and epilogueTileN to match tileM and tileN respectively"); + if (updateOptions) { + options.mEpilogueTileM = options.mTileM; + options.mEpilogueTileN = options.mTileN; + } else { + return false; + } + } + } + if (options.mSplitK == SplitK::Dsmem) { + TLLM_CHECK_ERROR(options.mClusterDimZ == options.mNumSlicesForSplitK, + "CGA size must be equal to the number of slices in split-k"); + } + + // Maps numStagesMma to (stagesWithinWorkTile, stagesAcrossWorkTile) if not already set. + // If (-1, -1) -> (numStagesMma / min(2, numStagesMma), min(2, numStagesMma)) + // If ( m, -1) -> (m, numStagesMma / m) + // If (-1, n) -> (numStagesMma / n, n) + if (options.mNumStagesMmaWithinWorkTile == -1 && options.mNumStagesMmaAcrossWorkTile == -1) { + if (updateOptions) { + options.mNumStagesMmaAcrossWorkTile = std::min(2, options.mNumStagesMma); + options.mNumStagesMmaWithinWorkTile = + options.mNumStagesMma / options.mNumStagesMmaAcrossWorkTile; + } else { + return false; + } + } else if (options.mNumStagesMmaWithinWorkTile == -1) { + if (updateOptions) { + options.mNumStagesMmaWithinWorkTile = + options.mNumStagesMma / options.mNumStagesMmaAcrossWorkTile; + } else { + return false; + } + } else if (options.mNumStagesMmaAcrossWorkTile == -1) { + if (updateOptions) { + options.mNumStagesMmaAcrossWorkTile = + options.mNumStagesMma / options.mNumStagesMmaWithinWorkTile; + } else { + return false; + } + } + // Check mma stages. + TLLM_CHECK_ERROR_FMT(options.mNumStagesMmaWithinWorkTile * options.mNumStagesMmaAcrossWorkTile == + options.mNumStagesMma && + options.mNumStagesMmaAcrossWorkTile <= 2, + "Condition numStagesMmaWithinWorkTile (%d) * numStagesMmaAcrossWorkTile " + "(%d) == numStagesMma (%d) && numStagesMmaAcrossWorkTile (%d) <= 2 must be " + "satisfied. Check arguments.", + options.mNumStagesMmaWithinWorkTile, options.mNumStagesMmaAcrossWorkTile, + options.mNumStagesMma, options.mNumStagesMmaAcrossWorkTile); + // Mma stage must be 1 for pre-Hopper. + TLLM_CHECK_ERROR(isBlackwell || options.mNumStagesMma == 1, + "Mma stage must be 1 for pre-Hopper. Found ", options.mNumStagesMma); + // DeepSeek Fp8 + if (!options.mUseDeepSeekFp8) { + TLLM_CHECK_ERROR(options.mNumStagesMmaWithinWorkTile == 1, + "Non-DeepSeekFp8 requires numStagesMmaWithinWorkTile == 1"); + if (options.mNumStagesMma > 1) { + TLLM_CHECK_ERROR(options.mTileScheduler == TileScheduler::Persistent, + "Non-DeepSeekFp8 requires persistent scheduler when using numStagesMma >1"); + } + } + if (options.mUseDeepSeekFp8) { + TLLM_CHECK_ERROR(options.mDtypeA == tg::Dtype::E4m3 && options.mDtypeB == tg::Dtype::E4m3, + "A and B dtype must be E4m3 for DeepSeek Fp8. Found dtypeA=", + tg::dtypeToString(options.mDtypeA), + " dtypeB=", tg::dtypeToString(options.mDtypeB)); + + TLLM_CHECK_ERROR(isBlackwell, "DeepSeek Fp8 is not supported for Hopper"); + TLLM_CHECK_ERROR(options.mAllReduceAlgo == AllReduceAlgo::None, + "DeepSeek Fp8 does not support AllReduce"); + + // Check that TileK = 128 for correct scaling of every 128 channels. + TLLM_CHECK_ERROR(options.mTileK == 128, "Tile-K must be equal to 128 for DeepSeek Fp8"); + TLLM_CHECK_ERROR(options.mK % options.mTileK == 0, "K must be a multiple of TileK"); + // Tile sizes of the output hidden dimension. + auto hiddenDimPerOutputTile = options.mTransposeMmaOutput ? options.mTileM : options.mTileN; + auto hiddenDimPerEpilogueTile = + options.mTransposeMmaOutput ? options.mEpilogueTileM : options.mEpilogueTileN; + auto hiddenDimPerMma = options.mTransposeMmaOutput ? options.mMmaM : options.mMmaN; + auto hiddenDimName = options.mTransposeMmaOutput ? "M" : "N"; + TLLM_CHECK_WARNING(options.mNumStagesMmaWithinWorkTile > 1, + "DeepSeekFp8 recommends setting \"-numStagesMmaWithinWorkTile 2\"."); + // Update the number of stages of the MMA accumulator pipeline. TODO: enable by default for + // deepseek. + // options.mNumStagesMma = 2; + // Use two MMA warps to reduce mbar trywait latency. TODO: enable by default for deepseek. + // options.mUseTwoMmaWarps = true; + + // Make sure the GEMM-K dimension is a multiple of 128 when using DeepSeek FP8. + TLLM_CHECK_ERROR(options.mK % 128 == 0, + "GEMM-K must be a multiple of 128 when using DeepSeek Fp8. Found ", + options.mK); + + // Check that the output tile N can be processed with the epilogue tile granularity. + TLLM_CHECK_ERROR((hiddenDimPerOutputTile / 2) % hiddenDimPerEpilogueTile == 0, + "DeepSeek Fp8 requires Tile", hiddenDimName, " / 2 (", + hiddenDimPerOutputTile / 2, ") being a multiple of EpilogueTile", + hiddenDimName, " (", hiddenDimPerEpilogueTile, ")"); + // Check that the output tile N can be processed with the epilogue tile granularity. + TLLM_CHECK_ERROR((hiddenDimPerOutputTile / 2) % hiddenDimPerMma == 0, + "DeepSeek Fp8 requires Tile", hiddenDimName, " / 2 (", + hiddenDimPerOutputTile / 2, ") being a multiple of mma", hiddenDimName, " (", + hiddenDimPerMma, ")"); + } + + if (options.mSliceK) { + TLLM_CHECK_ERROR(isBlackwell, "Slice-K is not supported on Hopper"); + + TLLM_CHECK_ERROR(!options.mUseDeepSeekFp8, "DeepSeek Fp8 GEMM is not supported for slice-K"); + TLLM_CHECK_ERROR(options.mUseTwoTmaLoadWarps, "Slice-K requires two warp load for A and B"); + TLLM_CHECK_ERROR(options.mTransposeMmaOutput, "Slice-K requires transpose mma output"); + TLLM_CHECK_ERROR(options.mUseShuffledMatrixA, "Slice-K requires shuffled matrix A"); + TLLM_CHECK_ERROR(options.mTileK % 128 == 0, "Slice-K requires TileK be a multiple of 128"); + TLLM_CHECK_ERROR(options.mMmaM == 128, "Slice-K requires MmaM == 128"); + TLLM_CHECK_ERROR(options.mTileN == options.mEpilogueTileN, + "TileN must be equal to EpilogueTileN for slice-K"); + + TLLM_LOG_WARNING("Overwriting TileM and EpilogueTileM to 32 for slice-K"); + if (options.mTileM != 32 || options.mEpilogueTileM != 32) { + if (updateOptions) { + // FIXME: it is possible to remove this restriction. + options.mTileM = 32; + options.mEpilogueTileM = 32; + } else { + return false; + } + } + TLLM_CHECK_ERROR(options.mDtypeA == tg::Dtype::E4m3 && options.mDtypeB == tg::Dtype::E4m3, + "Slice-K requires e4m3 input dtype"); + + if (options.mNumSlicesForSliceK != 4) { + if (updateOptions) { + options.mNumSlicesForSliceK = 4; + } else { + return false; + } + } + TLLM_CHECK_ERROR((options.mTileK / options.mMmaK) % options.mNumSlicesForSliceK == 0, "TileK (", + options.mTileK, ") / MmaK (", options.mMmaK, + ") must be a multiple of mNumSlicesForSliceK (", options.mNumSlicesForSliceK, + ")"); + } + + if (options.mUseUnrollLoop2xForMma) { + // Number of iterations in K dimension after padding. + // Note the perCtaK in each CTA in the splitK group are padded to the same number of iterations. + // E.g., K = 512, TileK = 128, numSlicesForSplitK = 3. Then the padded K is + // + // ceil(512 / (128*3)) * (128*3) = 768 + // + int paddedK = divUpMul(options.mK, options.mTileK * options.mNumSlicesForSplitK); + // Check that the padded K (K rounded to next multiple of tileK) is a multiple of 2*TileK when + // UnrollLoop2x is enabled. This is to avoid deadlock when mma runs even-numbered loop while the + // other warps run odd-numbered loop. + // + bool notSupported = (paddedK / options.mNumSlicesForSplitK) % (options.mTileK * 2) != 0; + if (notSupported) { + TLLM_LOG_WARNING("Size K / splitK must be a multiple of TileK * 2. Found TileK=", + options.mTileK, " and K=", options.mK, " (paddedK=", paddedK, + ") and numSlicesForSplitK=", options.mNumSlicesForSplitK, + ". Disabling unrollLoop2xForMma."); + if (updateOptions) { + options.mUseUnrollLoop2xForMma = false; + } else { + return false; + } + } + } + + if (!isBlackwell && options.mTileScheduler == TileScheduler::Persistent) { + // TODO(anchengc): will be supported in upcoming MRs. + TLLM_LOG_WARNING("Persistent scheduling is not supported on Hopper. Using Static scheduling."); + if (updateOptions) { + options.mTileScheduler = TileScheduler::Static; + } else { + return false; + } + } + + if (options.mEnablesDelayedEarlyExit && options.mEnablesEarlyExit) { + TLLM_LOG_WARNING( + "Only one of early exit and delayed early exit should be enabled. Disabling " + "delayed early exit"); + if (updateOptions) { + options.mEnablesDelayedEarlyExit = false; + } else { + return false; + } + } + + // This check prevents the triggering of the secondary (PREEXIT) from executing before the wait + // for primary (ACQBULK). This could lead to the following confusing situation, which we want to + // avoid: + // + // Kernel 3 is written with the assumption that it can read the output of + // kernel 1 *without* ACQBULK and the output of kernel 2 *with* ACQBULK. + // However, when we allow PREEXIT and ACQBULK to be executed out of order, + // this is not guaranteed. + // + // Time: ----> + // + // Kernel 1: ----PREEXIT-----------FLUSH + // Kernel 2: -------PREEXIT----ACQBULK---FLUSH + // Kernel 3: Warp 0: ---- (!) Output of 1,2 is not yet visible + // ----------------------- + // Warp 1: ---- (!) We normally assume that 1 is visible is not yet + // visible- Warp 2: -------------------ACQBULK-- Kernel 1,2 output visible + // ---------- + TLLM_CHECK_ERROR( + (options.mGridWaitForPrimaryA || !options.mGridTriggerSecondaryA), + "A: If a task triggers a secondary kernel, it must also wait for primary kernel."); + TLLM_CHECK_ERROR( + (options.mGridWaitForPrimaryB || !options.mGridTriggerSecondaryB), + "B: If a task triggers a secondary kernel, it must also wait for primary kernel."); + + if (options.mUsePerTokenSfA || options.mUsePerTokenSfB) { + // Checks applicable to both MetaFP8 and RoutingScalesOnInput + TLLM_CHECK_ERROR(!options.mUseDeepSeekFp8, + "DeepSeek FP8 and per-token scaling are not compatible"); + TLLM_CHECK_ERROR(isBlackwell, "Per-token scaling is not supported for Hopper"); + if (options.mUsePerTokenSfA && options.mUsePerTokenSfB) { + // MetaFP8 case + TLLM_CHECK_ERROR(options.mDtypeA == tg::Dtype::E4m3 && options.mDtypeB == tg::Dtype::E4m3, + "A and B dtype must be E4m3 for Meta Fp8. Found dtypeA=", + tg::dtypeToString(options.mDtypeA), + " dtypeB=", tg::dtypeToString(options.mDtypeB)); + } else { + // RoutingScalesOnInput case + TLLM_CHECK_ERROR((options.mUsePerTokenSfA && !options.mTransposeMmaOutput) || + (options.mUsePerTokenSfB && options.mTransposeMmaOutput), + "In RoutingScalesOnInput mode, perToken scales must be used on activations"); + } + } + + // The generation should support non K-major layouts for both A and B; however, it is unclear if + // there is a use-case + TLLM_CHECK_ERROR( + (options.mLayoutA == MatrixLayout::MajorK) || (options.mLayoutB == MatrixLayout::MajorK), + "At least one matrix must be in k-major layout"); + + // Some features are currently only support when both matrices are in K-major format + if (options.mLayoutB != MatrixLayout::MajorK || options.mLayoutB != MatrixLayout::MajorK) { + TLLM_CHECK_ERROR(isBlackwell, "Non K-major layouts are only supported on Blackwell"); + TLLM_CHECK_ERROR(options.mSplitK == SplitK::None, "Non K-major layouts do not support split K"); + } + if (options.mLayoutA == MatrixLayout::MajorMn) { + TLLM_CHECK_ERROR(tg::dtypeGetNumBits(options.mDtypeA) >= 8, + "Subbyte types only support K major layout"); + } + if (options.mLayoutB == MatrixLayout::MajorMn) { + TLLM_CHECK_ERROR(tg::dtypeGetNumBits(options.mDtypeB) >= 8, + "Subbyte types only support K major layout"); + } + + if ((options.mLayoutA == MatrixLayout::BlockMajorK) || + (options.mLayoutB == MatrixLayout::BlockMajorK)) { + bool const isBlockA = options.mLayoutA == MatrixLayout::BlockMajorK; + + // Block K size must be 128B. + // TODO Leaving this as an option for now in case we want to expertiment with other block sizes + // As the user is not expected to set this, do not fail if updateOptions is false + int32_t const elemSizeInBits = + (isBlockA) ? tg::dtypeGetNumBits(options.mDtypeA) : tg::dtypeGetNumBits(options.mDtypeB); + int32_t const elemsIn128B = 128 * 8 /* Bits in byte */ / elemSizeInBits; + + if (options.mBlockK != elemsIn128B) { + if (updateOptions) { + options.mBlockK = elemsIn128B; + } else { + return false; + } + } + + if (options.mBlockK > options.mTileK) { + TLLM_CHECK_ERROR( + options.mBlockK % options.mTileK == 0, + "If block size is greater than tile size, block size must be a multiple of tile size"); + } else if (options.mBlockK < options.mTileK) { + TLLM_CHECK_ERROR( + options.mTileK % options.mBlockK == 0, + "If tile size is greater than block size, tile size must be a multiple of block size"); + } + } + + if (!isBiasTypeNone(options.mBiasType)) { + TLLM_CHECK_ERROR(!isBiasTypeMn(options.mBiasType), "BiasType::Mn is not supported"); + TLLM_CHECK_ERROR(!options.mUseDeepSeekFp8, "Bias is not supported for DeepSeek Fp8"); + TLLM_CHECK_ERROR(!(options.mUsePerTokenSfA && options.mUsePerTokenSfB), + "Bias is not supported for Meta Fp8"); + } + + if (updateOptions) { + // Init kernel traits. + options.mKernelTraits = KernelTraits( + options.mDtypeA, options.mDtypeB, options.mDtypeC, options.mDtypeAcc, options.mDtypeMmaA, + options.mDtypeMmaB, options.mMmaKind, options.mTileM, options.mTileN, options.mTileK, + options.mEpilogueTileM, options.mEpilogueTileN, options.mNumStages, options.mNumStagesMma, + options.mNumSlicesForSplitK, options.mNumSlicesForSliceK, options.mSplitK, + options.mUseTmaStore, options.mTransposeMmaOutput, options.mAllReduceAlgo, + options.mTileScheduler == TileScheduler::Persistent, options.mUseDeepSeekFp8, + options.mUsePerTokenSfA, options.mUsePerTokenSfB, options.mBiasType); + } + + return true; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace gemm + +#ifdef TLLM_GEN_EXPORT_INTERFACE + +#undef TLLM_CHECK_ERROR +#undef TLLM_CHECK_ERROR_FMT +#undef TLLM_CHECK_WARNING +#undef TLLM_LOG_WARNING +#undef TLLM_LOG_INFO +#undef TLLM_LOG_ERROR + +#endif // TLLM_GEN_EXPORT_INTERFACE + +} // namespace batchedGemm diff --git a/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/KernelMetaInfo.h b/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/KernelMetaInfo.h new file mode 100644 index 0000000000000000000000000000000000000000..f5c064af71d0ecf2131c6d01d84502b380508b7b --- /dev/null +++ b/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/KernelMetaInfo.h @@ -0,0 +1,34727 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & + * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "BatchedGemmOptions.h" + +namespace batchedGemm { + +namespace tensorrt_llm { +namespace kernels { +// clang-format off + +#define TLLM_GEN_COMMIT "32110eb" +#define TLLM_GEN_EXPORT_VERSION "7.0.3.0.3.0" + +static constexpr size_t tllmGenBatchedGemmListLen = 408; + +#ifndef EXCLUDE_SM_100 +extern unsigned char* Bmm_Bfloat16_E2m1E2m1_Fp32_t128x16x256_s6_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E2m1E2m1_Fp32_t128x16x256_s6_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E2m1E2m1_Fp32_t128x16x256u2_s6_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E2m1E2m1_Fp32_t128x16x256u2_s6_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E2m1E2m1_Fp32_t128x16x512_s4_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E2m1E2m1_Fp32_t128x16x512_s4_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E2m1E2m1_Fp32_t128x16x512u2_s4_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E2m1E2m1_Fp32_t128x16x512u2_s4_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E2m1E2m1_Fp32_t128x32x256_s6_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E2m1E2m1_Fp32_t128x32x256_s6_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E2m1E2m1_Fp32_t128x32x256u2_s6_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E2m1E2m1_Fp32_t128x32x256u2_s6_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E2m1E2m1_Fp32_t128x32x512_s4_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E2m1E2m1_Fp32_t128x32x512_s4_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E2m1E2m1_Fp32_t128x32x512u2_s4_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E2m1E2m1_Fp32_t128x32x512u2_s4_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E2m1E2m1_Fp32_t128x64x256_s6_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E2m1E2m1_Fp32_t128x64x256_s6_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E2m1E2m1_Fp32_t128x64x256u2_s6_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E2m1E2m1_Fp32_t128x64x256u2_s6_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E2m1E2m1_Fp32_t128x64x512_s3_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E2m1E2m1_Fp32_t128x64x512_s3_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E2m1E2m1_Fp32_t128x64x512u2_s3_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E2m1E2m1_Fp32_t128x64x512u2_s3_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E2m1E2m1_Fp32_t128x8x256_s6_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E2m1E2m1_Fp32_t128x8x256_s6_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E2m1E2m1_Fp32_t128x8x256u2_s6_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E2m1E2m1_Fp32_t128x8x256u2_s6_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E2m1E2m1_Fp32_t128x8x512_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E2m1E2m1_Fp32_t128x8x512_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E2m1E2m1_Fp32_t128x8x512_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E2m1E2m1_Fp32_t128x8x512u2_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E2m1E2m1_Fp32_t128x8x512u2_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E2m1E2m1_Fp32_t128x8x512u2_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x16x256_s3_et128x16_m128x16x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x16x256_s3_et128x16_m128x16x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x16x256_s3_et128x16_m128x16x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x16x256_s3_et128x16_m128x16x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x16x256u2_s3_et128x16_m128x16x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x16x256u2_s3_et128x16_m128x16x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x16x256u2_s3_et128x16_m128x16x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x16x256u2_s3_et128x16_m128x16x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x32x256_s3_et128x32_m128x32x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x32x256_s3_et128x32_m128x32x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x32x256_s3_et128x32_m128x32x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x32x256_s3_et128x32_m128x32x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x32x256u2_s3_et128x32_m128x32x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x32x256u2_s3_et128x32_m128x32x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x32x256u2_s3_et128x32_m128x32x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x32x256u2_s3_et128x32_m128x32x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x64x256_s3_et128x64_m128x64x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x64x256_s3_et128x64_m128x64x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x64x256_s3_et128x64_m128x64x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x64x256_s3_et128x64_m128x64x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x64x256u2_s3_et128x64_m128x64x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x64x256u2_s3_et128x64_m128x64x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x64x256u2_s3_et128x64_m128x64x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x64x256u2_s3_et128x64_m128x64x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x8x256_s3_et128x8_m128x8x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x8x256_s3_et128x8_m128x8x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x8x256_s3_et128x8_m128x8x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x8x256_s3_et128x8_m128x8x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x8x256u2_s3_et128x8_m128x8x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x8x256u2_s3_et128x8_m128x8x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x8x256u2_s3_et128x8_m128x8x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x8x256u2_s3_et128x8_m128x8x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x16x256_s5_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x16x256_s5_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x16x256u2_s5_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x16x256u2_s5_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x32x256_s5_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x32x256_s5_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x32x256u2_s5_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x32x256u2_s5_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x64x256_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x64x256_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x64x256u2_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x64x256u2_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x16x256_s3_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x16x256_s3_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x16x256_s4_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x16x256_s4_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x16x256u2_s3_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x16x256u2_s3_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x16x256u2_s4_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x16x256u2_s4_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x32x256_s3_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x32x256_s3_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x32x256_s4_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x32x256_s4_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x32x256u2_s3_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x32x256u2_s3_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x32x256u2_s4_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x32x256u2_s4_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x64x256_s3_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x64x256_s3_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x64x256_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x64x256_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x64x256u2_s3_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x64x256u2_s3_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x64x256u2_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x64x256u2_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s4_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s4_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s4_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s4_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s6_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s6_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s6_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s6_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s4_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s4_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s4_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s4_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s6_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s6_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s6_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s6_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E2m1_E2m1E2m1_Fp32_t128x16x512_s4_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E2m1_E2m1E2m1_Fp32_t128x16x512_s4_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E2m1_E2m1E2m1_Fp32_t128x16x512u2_s4_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E2m1_E2m1E2m1_Fp32_t128x16x512u2_s4_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E2m1_E2m1E2m1_Fp32_t128x32x512_s4_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E2m1_E2m1E2m1_Fp32_t128x32x512_s4_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E2m1_E2m1E2m1_Fp32_t128x32x512u2_s4_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E2m1_E2m1E2m1_Fp32_t128x32x512u2_s4_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E2m1_E2m1E2m1_Fp32_t128x64x256_s6_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E2m1_E2m1E2m1_Fp32_t128x64x256_s6_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E2m1_E2m1E2m1_Fp32_t128x64x256u2_s6_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E2m1_E2m1E2m1_Fp32_t128x64x256u2_s6_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E2m1_E2m1E2m1_Fp32_t128x8x512_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E2m1_E2m1E2m1_Fp32_t128x8x512_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E2m1_E2m1E2m1_Fp32_t128x8x512_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E2m1_E2m1E2m1_Fp32_t128x8x512u2_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E2m1_E2m1E2m1_Fp32_t128x8x512u2_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E2m1_E2m1E2m1_Fp32_t128x8x512u2_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x16x512_s3_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_tokSfB_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x16x512u2_s3_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_tokSfB_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x32x512_s2_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_tokSfB_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x32x512u2_s2_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_tokSfB_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x64x512_s2_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_tokSfB_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x64x512u2_s2_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_tokSfB_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_tokSfB_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_E4m3E4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_tokSfB_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x16x256_s5_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x16x256_s5_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x16x256u2_s5_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x16x256u2_s5_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x32x256_s5_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x32x256_s5_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x32x256u2_s5_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x32x256u2_s5_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x64x256_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x64x256_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x64x256u2_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x64x256u2_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Fp16_E2m1E2m1_Fp32_t128x8x512_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Fp16_E2m1E2m1_Fp32_t128x8x512u2_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Fp16_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Fp16_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Fp16_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Fp16_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Fp16_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Fp16_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Fp16_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Fp16_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Fp16_E4m3E4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a_cubin = nullptr; +extern unsigned char* Bmm_Fp16_E4m3E4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x16x256_s3_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x16x256_s3_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x16x256_s4_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x16x256_s4_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x16x256u2_s3_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x16x256u2_s3_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x16x256u2_s4_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x16x256u2_s4_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x32x256_s3_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x32x256_s3_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x32x256_s4_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x32x256_s4_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x32x256u2_s3_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x32x256u2_s3_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x32x256u2_s4_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x32x256u2_s4_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x64x256_s3_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x64x256_s3_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x64x256_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x64x256_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x64x256u2_s3_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x64x256u2_s3_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x64x256u2_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x64x256u2_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s4_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s4_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s4_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s4_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s6_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s6_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s6_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s6_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s4_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s4_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s4_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s4_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s6_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s6_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s6_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s6_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +extern unsigned char* Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin = nullptr; +#endif // EXCLUDE_SM_100 + +#ifndef EXCLUDE_SM_100 +inline unsigned int Bmm_Bfloat16_E2m1E2m1_Fp32_t128x16x256_s6_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E2m1E2m1_Fp32_t128x16x256_s6_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E2m1E2m1_Fp32_t128x16x256u2_s6_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E2m1E2m1_Fp32_t128x16x256u2_s6_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E2m1E2m1_Fp32_t128x16x512_s4_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E2m1E2m1_Fp32_t128x16x512_s4_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E2m1E2m1_Fp32_t128x16x512u2_s4_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E2m1E2m1_Fp32_t128x16x512u2_s4_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E2m1E2m1_Fp32_t128x32x256_s6_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E2m1E2m1_Fp32_t128x32x256_s6_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E2m1E2m1_Fp32_t128x32x256u2_s6_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E2m1E2m1_Fp32_t128x32x256u2_s6_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E2m1E2m1_Fp32_t128x32x512_s4_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E2m1E2m1_Fp32_t128x32x512_s4_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E2m1E2m1_Fp32_t128x32x512u2_s4_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E2m1E2m1_Fp32_t128x32x512u2_s4_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E2m1E2m1_Fp32_t128x64x256_s6_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E2m1E2m1_Fp32_t128x64x256_s6_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E2m1E2m1_Fp32_t128x64x256u2_s6_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E2m1E2m1_Fp32_t128x64x256u2_s6_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E2m1E2m1_Fp32_t128x64x512_s3_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E2m1E2m1_Fp32_t128x64x512_s3_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E2m1E2m1_Fp32_t128x64x512u2_s3_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E2m1E2m1_Fp32_t128x64x512u2_s3_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E2m1E2m1_Fp32_t128x8x256_s6_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E2m1E2m1_Fp32_t128x8x256_s6_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E2m1E2m1_Fp32_t128x8x256u2_s6_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E2m1E2m1_Fp32_t128x8x256u2_s6_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E2m1E2m1_Fp32_t128x8x512_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E2m1E2m1_Fp32_t128x8x512_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E2m1E2m1_Fp32_t128x8x512_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E2m1E2m1_Fp32_t128x8x512u2_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E2m1E2m1_Fp32_t128x8x512u2_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E2m1E2m1_Fp32_t128x8x512u2_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x16x256_s3_et128x16_m128x16x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x16x256_s3_et128x16_m128x16x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x16x256_s3_et128x16_m128x16x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x16x256_s3_et128x16_m128x16x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x16x256u2_s3_et128x16_m128x16x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x16x256u2_s3_et128x16_m128x16x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x16x256u2_s3_et128x16_m128x16x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x16x256u2_s3_et128x16_m128x16x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x32x256_s3_et128x32_m128x32x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x32x256_s3_et128x32_m128x32x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x32x256_s3_et128x32_m128x32x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x32x256_s3_et128x32_m128x32x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x32x256u2_s3_et128x32_m128x32x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x32x256u2_s3_et128x32_m128x32x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x32x256u2_s3_et128x32_m128x32x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x32x256u2_s3_et128x32_m128x32x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x64x256_s3_et128x64_m128x64x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x64x256_s3_et128x64_m128x64x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x64x256_s3_et128x64_m128x64x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x64x256_s3_et128x64_m128x64x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x64x256u2_s3_et128x64_m128x64x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x64x256u2_s3_et128x64_m128x64x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x64x256u2_s3_et128x64_m128x64x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x64x256u2_s3_et128x64_m128x64x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x8x256_s3_et128x8_m128x8x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x8x256_s3_et128x8_m128x8x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x8x256_s3_et128x8_m128x8x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x8x256_s3_et128x8_m128x8x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x8x256u2_s3_et128x8_m128x8x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x8x256u2_s3_et128x8_m128x8x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x8x256u2_s3_et128x8_m128x8x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x8x256u2_s3_et128x8_m128x8x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x16x256_s5_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x16x256_s5_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x16x256u2_s5_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x16x256u2_s5_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x32x256_s5_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x32x256_s5_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x32x256u2_s5_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x32x256u2_s5_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x64x256_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x64x256_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x64x256u2_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x64x256u2_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x16x256_s3_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x16x256_s3_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x16x256_s4_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x16x256_s4_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x16x256u2_s3_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x16x256u2_s3_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x16x256u2_s4_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x16x256u2_s4_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x32x256_s3_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x32x256_s3_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x32x256_s4_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x32x256_s4_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x32x256u2_s3_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x32x256u2_s3_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x32x256u2_s4_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x32x256u2_s4_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x64x256_s3_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x64x256_s3_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x64x256_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x64x256_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x64x256u2_s3_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x64x256u2_s3_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x64x256u2_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x64x256u2_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s4_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s4_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s4_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s4_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s6_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s6_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s6_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s6_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s4_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s4_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s4_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s4_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s6_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s6_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s6_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s6_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E2m1_E2m1E2m1_Fp32_t128x16x512_s4_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E2m1_E2m1E2m1_Fp32_t128x16x512_s4_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E2m1_E2m1E2m1_Fp32_t128x16x512u2_s4_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E2m1_E2m1E2m1_Fp32_t128x16x512u2_s4_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E2m1_E2m1E2m1_Fp32_t128x32x512_s4_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E2m1_E2m1E2m1_Fp32_t128x32x512_s4_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E2m1_E2m1E2m1_Fp32_t128x32x512u2_s4_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E2m1_E2m1E2m1_Fp32_t128x32x512u2_s4_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E2m1_E2m1E2m1_Fp32_t128x64x256_s6_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E2m1_E2m1E2m1_Fp32_t128x64x256_s6_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E2m1_E2m1E2m1_Fp32_t128x64x256u2_s6_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E2m1_E2m1E2m1_Fp32_t128x64x256u2_s6_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E2m1_E2m1E2m1_Fp32_t128x8x512_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E2m1_E2m1E2m1_Fp32_t128x8x512_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a_cubin_len = 0; +inline unsigned int Bmm_E2m1_E2m1E2m1_Fp32_t128x8x512_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E2m1_E2m1E2m1_Fp32_t128x8x512u2_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E2m1_E2m1E2m1_Fp32_t128x8x512u2_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a_cubin_len = 0; +inline unsigned int Bmm_E2m1_E2m1E2m1_Fp32_t128x8x512u2_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x16x512_s3_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_tokSfB_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x16x512u2_s3_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_tokSfB_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x32x512_s2_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_tokSfB_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x32x512u2_s2_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_tokSfB_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x64x512_s2_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_tokSfB_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x64x512u2_s2_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_tokSfB_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_tokSfB_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_E4m3E4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_tokSfB_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x16x256_s5_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x16x256_s5_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x16x256u2_s5_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x16x256u2_s5_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x32x256_s5_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x32x256_s5_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x32x256u2_s5_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x32x256u2_s5_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x64x256_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x64x256_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x64x256u2_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x64x256u2_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_Fp16_E2m1E2m1_Fp32_t128x8x512_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a_cubin_len = 0; +inline unsigned int Bmm_Fp16_E2m1E2m1_Fp32_t128x8x512u2_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a_cubin_len = 0; +inline unsigned int Bmm_Fp16_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_sm100a_cubin_len = 0; +inline unsigned int Bmm_Fp16_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a_cubin_len = 0; +inline unsigned int Bmm_Fp16_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_sm100a_cubin_len = 0; +inline unsigned int Bmm_Fp16_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a_cubin_len = 0; +inline unsigned int Bmm_Fp16_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_sm100a_cubin_len = 0; +inline unsigned int Bmm_Fp16_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a_cubin_len = 0; +inline unsigned int Bmm_Fp16_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_sm100a_cubin_len = 0; +inline unsigned int Bmm_Fp16_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a_cubin_len = 0; +inline unsigned int Bmm_Fp16_E4m3E4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a_cubin_len = 0; +inline unsigned int Bmm_Fp16_E4m3E4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x16x256_s3_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x16x256_s3_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x16x256_s4_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x16x256_s4_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x16x256u2_s3_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x16x256u2_s3_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x16x256u2_s4_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x16x256u2_s4_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x32x256_s3_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x32x256_s3_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x32x256_s4_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x32x256_s4_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x32x256u2_s3_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x32x256u2_s3_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x32x256u2_s4_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x32x256u2_s4_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x64x256_s3_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x64x256_s3_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x64x256_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x64x256_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x64x256u2_s3_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x64x256u2_s3_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x64x256u2_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x64x256u2_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s4_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s4_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s4_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s4_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s6_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s6_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s6_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s6_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s4_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s4_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s4_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s4_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s6_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s6_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s6_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s6_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +inline unsigned int Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len = 0; +#endif // EXCLUDE_SM_100 + + +static const batchedGemm::BatchedGemmConfig tllmGenBatchedGemmList[] = { +#ifndef EXCLUDE_SM_100 +{Bmm_Bfloat16_E2m1E2m1_Fp32_t128x16x256_s6_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E2m1E2m1_Fp32_t128x16x256_s6_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin_len, 136192, "bmm_Bfloat16_E2m1E2m1_Fp32_t128x16x256_s6_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a", 480, "3f50040873bad87a6494c9cfe1fdb74cf105a262da64d6eeffb7cde948e4c8b0", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826818) +, /* mDtypeB */ trtllm::gen::Dtype(17826818) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17826818) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaKind */ trtllm::gen::MmaKind(4) +, /* mMmaM */ 128 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E2m1E2m1_Fp32_t128x16x256_s6_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E2m1E2m1_Fp32_t128x16x256_s6_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin_len, 136192, "bmm_Bfloat16_E2m1E2m1_Fp32_t128x16x256_s6_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a", 448, "421c9180d5ef84c582c13ca6f758e3eb88db2c0221d6d5358481994c793fff9e", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826818) +, /* mDtypeB */ trtllm::gen::Dtype(17826818) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17826818) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaKind */ trtllm::gen::MmaKind(4) +, /* mMmaM */ 128 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E2m1E2m1_Fp32_t128x16x256u2_s6_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E2m1E2m1_Fp32_t128x16x256u2_s6_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin_len, 136192, "bmm_Bfloat16_E2m1E2m1_Fp32_t128x16x256u2_s6_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a", 480, "9a5af1b5f22f8cfad3585bc94c7ab9e3889608184af7a90fa72355ec3580e37b", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826818) +, /* mDtypeB */ trtllm::gen::Dtype(17826818) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17826818) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaKind */ trtllm::gen::MmaKind(4) +, /* mMmaM */ 128 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E2m1E2m1_Fp32_t128x16x256u2_s6_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E2m1E2m1_Fp32_t128x16x256u2_s6_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin_len, 136192, "bmm_Bfloat16_E2m1E2m1_Fp32_t128x16x256u2_s6_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a", 448, "fbb8bd03a8e4ccf046a6d26ba70538b6c9295f38d9558fe74c472dc6f9500407", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826818) +, /* mDtypeB */ trtllm::gen::Dtype(17826818) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17826818) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaKind */ trtllm::gen::MmaKind(4) +, /* mMmaM */ 128 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E2m1E2m1_Fp32_t128x16x512_s4_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E2m1E2m1_Fp32_t128x16x512_s4_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin_len, 178176, "bmm_Bfloat16_E2m1E2m1_Fp32_t128x16x512_s4_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a", 480, "1211763de22658770e6fe0eeeaf075b6c3b9fc404184e27919e920c9cd89ae6d", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826818) +, /* mDtypeB */ trtllm::gen::Dtype(17826818) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17826818) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaKind */ trtllm::gen::MmaKind(4) +, /* mMmaM */ 128 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E2m1E2m1_Fp32_t128x16x512_s4_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E2m1E2m1_Fp32_t128x16x512_s4_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin_len, 178176, "bmm_Bfloat16_E2m1E2m1_Fp32_t128x16x512_s4_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a", 448, "80f36089bf16902e0cfe5530fdf73181181e42bbe1148a06f528c9fcfe951c1f", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826818) +, /* mDtypeB */ trtllm::gen::Dtype(17826818) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17826818) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaKind */ trtllm::gen::MmaKind(4) +, /* mMmaM */ 128 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E2m1E2m1_Fp32_t128x16x512u2_s4_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E2m1E2m1_Fp32_t128x16x512u2_s4_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin_len, 178176, "bmm_Bfloat16_E2m1E2m1_Fp32_t128x16x512u2_s4_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a", 480, "0382740f8904e76ac8ba65bd139c0703098be937ca6a48f8d805828d2e9fb4d9", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826818) +, /* mDtypeB */ trtllm::gen::Dtype(17826818) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17826818) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaKind */ trtllm::gen::MmaKind(4) +, /* mMmaM */ 128 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E2m1E2m1_Fp32_t128x16x512u2_s4_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E2m1E2m1_Fp32_t128x16x512u2_s4_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin_len, 178176, "bmm_Bfloat16_E2m1E2m1_Fp32_t128x16x512u2_s4_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a", 448, "e9ca5eae59b91497e401da818bda7a66c07c5c03b9b12870b8b46b1ead340290", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826818) +, /* mDtypeB */ trtllm::gen::Dtype(17826818) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17826818) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaKind */ trtllm::gen::MmaKind(4) +, /* mMmaM */ 128 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E2m1E2m1_Fp32_t128x32x256_s6_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E2m1E2m1_Fp32_t128x32x256_s6_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin_len, 154624, "bmm_Bfloat16_E2m1E2m1_Fp32_t128x32x256_s6_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a", 480, "a18ea12a8da554dd67a5cade11d3e2162bb2cca0349a3858d2b18de571423898", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826818) +, /* mDtypeB */ trtllm::gen::Dtype(17826818) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17826818) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaKind */ trtllm::gen::MmaKind(4) +, /* mMmaM */ 128 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E2m1E2m1_Fp32_t128x32x256_s6_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E2m1E2m1_Fp32_t128x32x256_s6_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin_len, 154624, "bmm_Bfloat16_E2m1E2m1_Fp32_t128x32x256_s6_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a", 448, "c8691567af794c5de14d5a5cc4f2a5d28536ca6073422b4df47c9afbb8f2be06", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826818) +, /* mDtypeB */ trtllm::gen::Dtype(17826818) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17826818) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaKind */ trtllm::gen::MmaKind(4) +, /* mMmaM */ 128 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E2m1E2m1_Fp32_t128x32x256u2_s6_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E2m1E2m1_Fp32_t128x32x256u2_s6_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin_len, 154624, "bmm_Bfloat16_E2m1E2m1_Fp32_t128x32x256u2_s6_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a", 480, "28d8f178f4347d48c162b4a35786bdab9378c9dc471a02c80b3aed6650a4e6a5", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826818) +, /* mDtypeB */ trtllm::gen::Dtype(17826818) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17826818) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaKind */ trtllm::gen::MmaKind(4) +, /* mMmaM */ 128 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E2m1E2m1_Fp32_t128x32x256u2_s6_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E2m1E2m1_Fp32_t128x32x256u2_s6_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin_len, 154624, "bmm_Bfloat16_E2m1E2m1_Fp32_t128x32x256u2_s6_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a", 448, "b8269703cdd5409fa0bad9f8ed695dca41ff739af1ef94b346c3f0be0c404760", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826818) +, /* mDtypeB */ trtllm::gen::Dtype(17826818) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17826818) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaKind */ trtllm::gen::MmaKind(4) +, /* mMmaM */ 128 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E2m1E2m1_Fp32_t128x32x512_s4_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E2m1E2m1_Fp32_t128x32x512_s4_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin_len, 200704, "bmm_Bfloat16_E2m1E2m1_Fp32_t128x32x512_s4_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a", 480, "f30f8d237cd1cf14e5ec6e46384e15fb0f29b49ac06d115760f7cd5de96139d3", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826818) +, /* mDtypeB */ trtllm::gen::Dtype(17826818) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17826818) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaKind */ trtllm::gen::MmaKind(4) +, /* mMmaM */ 128 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E2m1E2m1_Fp32_t128x32x512_s4_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E2m1E2m1_Fp32_t128x32x512_s4_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin_len, 200704, "bmm_Bfloat16_E2m1E2m1_Fp32_t128x32x512_s4_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a", 448, "9bf99980b7a6c15a6ae98733aeec0db642207ff052c875884147723f39b6413b", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826818) +, /* mDtypeB */ trtllm::gen::Dtype(17826818) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17826818) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaKind */ trtllm::gen::MmaKind(4) +, /* mMmaM */ 128 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E2m1E2m1_Fp32_t128x32x512u2_s4_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E2m1E2m1_Fp32_t128x32x512u2_s4_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin_len, 200704, "bmm_Bfloat16_E2m1E2m1_Fp32_t128x32x512u2_s4_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a", 480, "78318af7a0405b458217987427d2fe31e668b1e5316c3d14d72ec23b71117aaa", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826818) +, /* mDtypeB */ trtllm::gen::Dtype(17826818) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17826818) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaKind */ trtllm::gen::MmaKind(4) +, /* mMmaM */ 128 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E2m1E2m1_Fp32_t128x32x512u2_s4_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E2m1E2m1_Fp32_t128x32x512u2_s4_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin_len, 200704, "bmm_Bfloat16_E2m1E2m1_Fp32_t128x32x512u2_s4_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a", 448, "c9df13716958e4441c1d29ac5c1fa6cb2d293389fd8f779bc03fd2abfe62423a", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826818) +, /* mDtypeB */ trtllm::gen::Dtype(17826818) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17826818) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaKind */ trtllm::gen::MmaKind(4) +, /* mMmaM */ 128 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E2m1E2m1_Fp32_t128x64x256_s6_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E2m1E2m1_Fp32_t128x64x256_s6_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin_len, 190464, "bmm_Bfloat16_E2m1E2m1_Fp32_t128x64x256_s6_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a", 480, "0a9146fb80f1248e9b89895999cb31834c3994ea3b490ad9a8a444386742d16f", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826818) +, /* mDtypeB */ trtllm::gen::Dtype(17826818) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17826818) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaKind */ trtllm::gen::MmaKind(4) +, /* mMmaM */ 128 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E2m1E2m1_Fp32_t128x64x256_s6_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E2m1E2m1_Fp32_t128x64x256_s6_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin_len, 190464, "bmm_Bfloat16_E2m1E2m1_Fp32_t128x64x256_s6_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a", 448, "bd37ba431b1765a9ddf90d5f46549ef5bd453a20037cdd497baa79be476b781b", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826818) +, /* mDtypeB */ trtllm::gen::Dtype(17826818) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17826818) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaKind */ trtllm::gen::MmaKind(4) +, /* mMmaM */ 128 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E2m1E2m1_Fp32_t128x64x256u2_s6_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E2m1E2m1_Fp32_t128x64x256u2_s6_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin_len, 190464, "bmm_Bfloat16_E2m1E2m1_Fp32_t128x64x256u2_s6_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a", 480, "f6435ef5c7d9d09aed443595c849477def79f95a9995dc0a938aed6ff6591a16", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826818) +, /* mDtypeB */ trtllm::gen::Dtype(17826818) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17826818) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaKind */ trtllm::gen::MmaKind(4) +, /* mMmaM */ 128 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E2m1E2m1_Fp32_t128x64x256u2_s6_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E2m1E2m1_Fp32_t128x64x256u2_s6_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin_len, 190464, "bmm_Bfloat16_E2m1E2m1_Fp32_t128x64x256u2_s6_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a", 448, "47277e74746b699eb163ed30cb063fec743743f969b881247a8723296927a452", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826818) +, /* mDtypeB */ trtllm::gen::Dtype(17826818) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17826818) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaKind */ trtllm::gen::MmaKind(4) +, /* mMmaM */ 128 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E2m1E2m1_Fp32_t128x64x512_s3_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E2m1E2m1_Fp32_t128x64x512_s3_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin_len, 190464, "bmm_Bfloat16_E2m1E2m1_Fp32_t128x64x512_s3_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a", 480, "4226a60504e7275e4b4a42f764e98294258e66f8cc4a10e14e171df7b2f5066b", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826818) +, /* mDtypeB */ trtllm::gen::Dtype(17826818) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17826818) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaKind */ trtllm::gen::MmaKind(4) +, /* mMmaM */ 128 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E2m1E2m1_Fp32_t128x64x512_s3_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E2m1E2m1_Fp32_t128x64x512_s3_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin_len, 190464, "bmm_Bfloat16_E2m1E2m1_Fp32_t128x64x512_s3_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a", 448, "bf4fafd9e3cb30a9b0446dd8838849db06a959f16689df21dfad17513682bc63", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826818) +, /* mDtypeB */ trtllm::gen::Dtype(17826818) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17826818) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaKind */ trtllm::gen::MmaKind(4) +, /* mMmaM */ 128 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E2m1E2m1_Fp32_t128x64x512u2_s3_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E2m1E2m1_Fp32_t128x64x512u2_s3_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin_len, 190464, "bmm_Bfloat16_E2m1E2m1_Fp32_t128x64x512u2_s3_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a", 480, "86f8fffb1cea50303e40100d3bc51e0be57b3a70adc2a1859a8694f79049349d", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826818) +, /* mDtypeB */ trtllm::gen::Dtype(17826818) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17826818) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaKind */ trtllm::gen::MmaKind(4) +, /* mMmaM */ 128 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E2m1E2m1_Fp32_t128x64x512u2_s3_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E2m1E2m1_Fp32_t128x64x512u2_s3_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin_len, 190464, "bmm_Bfloat16_E2m1E2m1_Fp32_t128x64x512u2_s3_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a", 448, "1e121aa979c8bd4768d793bdb0003889e7be3f4a31ff2624b285fe4c92e7cb87", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826818) +, /* mDtypeB */ trtllm::gen::Dtype(17826818) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17826818) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaKind */ trtllm::gen::MmaKind(4) +, /* mMmaM */ 128 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E2m1E2m1_Fp32_t128x8x256_s6_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E2m1E2m1_Fp32_t128x8x256_s6_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin_len, 126976, "bmm_Bfloat16_E2m1E2m1_Fp32_t128x8x256_s6_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a", 480, "7a0b234601ff19b9a74f9a5e049cfa9bbb2412a588027cabb73ac41726388a05", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826818) +, /* mDtypeB */ trtllm::gen::Dtype(17826818) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17826818) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaKind */ trtllm::gen::MmaKind(4) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E2m1E2m1_Fp32_t128x8x256_s6_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E2m1E2m1_Fp32_t128x8x256_s6_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin_len, 126976, "bmm_Bfloat16_E2m1E2m1_Fp32_t128x8x256_s6_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a", 448, "273bd4ae9a7fcd4dc3b4beacae4660a8aeafe93ec01cf13c99ddda7cde90c6f6", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826818) +, /* mDtypeB */ trtllm::gen::Dtype(17826818) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17826818) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaKind */ trtllm::gen::MmaKind(4) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E2m1E2m1_Fp32_t128x8x256u2_s6_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E2m1E2m1_Fp32_t128x8x256u2_s6_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin_len, 126976, "bmm_Bfloat16_E2m1E2m1_Fp32_t128x8x256u2_s6_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a", 480, "235c7cecc53e4b5b0fca2a72f8e9a467e9fee921739087acf3a27cdebc688cc2", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826818) +, /* mDtypeB */ trtllm::gen::Dtype(17826818) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17826818) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaKind */ trtllm::gen::MmaKind(4) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E2m1E2m1_Fp32_t128x8x256u2_s6_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E2m1E2m1_Fp32_t128x8x256u2_s6_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin_len, 126976, "bmm_Bfloat16_E2m1E2m1_Fp32_t128x8x256u2_s6_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a", 448, "f14ef8c6fb913f5fe433f4f7cef57d22453d9827e152b7998154f6701193d4c7", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826818) +, /* mDtypeB */ trtllm::gen::Dtype(17826818) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17826818) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaKind */ trtllm::gen::MmaKind(4) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E2m1E2m1_Fp32_t128x8x512_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E2m1E2m1_Fp32_t128x8x512_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin_len, 166912, "bmm_Bfloat16_E2m1E2m1_Fp32_t128x8x512_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a", 480, "37779c1a8bc4082b5c10f813cec67884a076d704450abb72006730c9d643da3c", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826818) +, /* mDtypeB */ trtllm::gen::Dtype(17826818) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17826818) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaKind */ trtllm::gen::MmaKind(4) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E2m1E2m1_Fp32_t128x8x512_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E2m1E2m1_Fp32_t128x8x512_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin_len, 166912, "bmm_Bfloat16_E2m1E2m1_Fp32_t128x8x512_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a", 448, "82c9d523368e936cecb66b5964259c128095eefddc099fa5679c963d70ae1a86", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826818) +, /* mDtypeB */ trtllm::gen::Dtype(17826818) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17826818) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaKind */ trtllm::gen::MmaKind(4) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E2m1E2m1_Fp32_t128x8x512_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a_cubin, Bmm_Bfloat16_E2m1E2m1_Fp32_t128x8x512_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a_cubin_len, 166912, "bmm_Bfloat16_E2m1E2m1_Fp32_t128x8x512_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a", 448, "132269771417a940aa39f43c0dbf496c585ad385b38dc182f28274c97307aca4", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826818) +, /* mDtypeB */ trtllm::gen::Dtype(17826818) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17826818) +, /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaKind */ trtllm::gen::MmaKind(4) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 1 +, /* mNumTokens */ 0 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E2m1E2m1_Fp32_t128x8x512u2_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E2m1E2m1_Fp32_t128x8x512u2_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin_len, 166912, "bmm_Bfloat16_E2m1E2m1_Fp32_t128x8x512u2_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a", 480, "f564fad6aed8f446292c1019b7ec6fd8700afefe5ae9db03f9338e153fc8685f", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826818) +, /* mDtypeB */ trtllm::gen::Dtype(17826818) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17826818) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaKind */ trtllm::gen::MmaKind(4) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E2m1E2m1_Fp32_t128x8x512u2_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E2m1E2m1_Fp32_t128x8x512u2_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin_len, 166912, "bmm_Bfloat16_E2m1E2m1_Fp32_t128x8x512u2_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a", 448, "298d1a1794a17c5d0585684e0b33711daaa2f6e0b00c8f47dc3f9e08f68c79f7", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826818) +, /* mDtypeB */ trtllm::gen::Dtype(17826818) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17826818) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaKind */ trtllm::gen::MmaKind(4) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E2m1E2m1_Fp32_t128x8x512u2_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a_cubin, Bmm_Bfloat16_E2m1E2m1_Fp32_t128x8x512u2_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a_cubin_len, 166912, "bmm_Bfloat16_E2m1E2m1_Fp32_t128x8x512u2_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a", 448, "7bb43ecdf9e503d1f0b510b6161ca4c7c98390fb36255fa2c4c094ed474cf381", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826818) +, /* mDtypeB */ trtllm::gen::Dtype(17826818) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17826818) +, /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaKind */ trtllm::gen::MmaKind(4) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 1 +, /* mNumTokens */ 0 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len, 121856, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a", 416, "033f3c0b194580bd4a8e0d32296ec950d3336a90f22c1d2c19601b3f04706025", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len, 121856, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a", 384, "a4853e328a5ecc75c50ffa766da7f3dfc2474bab0d31e5b17fca6323b6dd06b9", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len, 121856, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a", 416, "10f1d796023187047bba9a36b6b4be2cf3625804b39e9119ba60c8f9ebafea56", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len, 121856, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a", 384, "e90f00f9cc3a99474284d1d40300be041cc4616412dedfbf3af6fd8991fee6aa", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len, 121856, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a", 416, "b9ff246be6e1494a34a8fd737cad55ce0d7cd492f0ad0981c11e7d8e56a1b6fe", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len, 121856, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a", 384, "4c7c2cfaf1a4ff72cbcc6da48792ebf84c4a01cb17e076bf5bc71a4ca3863853", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len, 121856, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a", 416, "23af50bafd4ff3219b3b2e77185babd14dc564196e59d4a7ce0c71f9f63f6e1d", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len, 121856, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a", 384, "63ec7f175f16f2148f9112466e7759c6d080de0f21c0657ad4fe32903ed220e8", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len, 121856, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a", 416, "0227acf26857469604694c037b2ec247e95f1d363b502579227bddd36474d22f", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len, 121856, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a", 384, "7649909f8952a983e3573f19be8d0eeef90ae853f3730cae014dd97bfe3ca225", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len, 121856, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a", 416, "0b0535fe5e1dc64350dc3b525ae1cfa4edf7c8e5ce140431c474786041e74f53", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len, 121856, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a", 384, "f40181a72c1a3a9de321ea65d10398214ee7f7de27a8e74e67b6d32e58194807", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len, 121856, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a", 416, "87fe035315d329d83ab5dcd4201a173a43190ee7ee2949a0836309afad140734", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len, 121856, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a", 384, "c92032ee7452c8f79ac9526fb9db84b0d50ab18da19011dfdcba6e6b3d852397", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len, 121856, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a", 416, "033db21cba9333e6789784c5d851b60e431e5d421753bd2713cf3da0063fc283", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len, 121856, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a", 384, "3be81e7035ee5a5d3488f9501197fea2caf7159de69217eab8d041dd013f5f7c", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len, 97280, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a", 416, "97266c0137ed61bca912f69fb299aa6369fbb4bc53ce3883031308e93c8923df", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len, 97280, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a", 384, "bf85855e9e7a943c409e87a393fab52c09d8fea4e3d49e750f8b463576555d23", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len, 97280, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a", 416, "e749eb37196ae08cf93ea4f2a7e65b085c2cad49d109102ee211748953e57b01", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len, 97280, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a", 384, "4cd14db05dde5b92cfff810910d8fca58b1607850cba9a6fcaac41dc204398e5", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len, 97280, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a", 416, "834628f9c031e0e74b5634a732bb1af96039f7736fb2c39c29efe97a60435d80", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len, 97280, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a", 384, "c6f35190ec003d65f903d0d017577377ffe50a4c9ba8ef7398232e514e775bcd", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len, 97280, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a", 416, "7400ee66aa0eaa3218c7e43e06c3b449bd9d263fee808a62e3019cd67cf71c78", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len, 97280, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a", 384, "401fd2bfcd51e41ed9abd426842f6d9a8e59ef3b85fff0b4e43a15a59dce992a", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len, 97280, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a", 416, "dc8fa6335b2667f779f474cbb5cfe8931fdc8b6ea6b54bce03d15f34a555b71c", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len, 97280, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a", 384, "e2e6c0e37a9bc30aedd1219d2e907f25d70ed2d0c507e82b4eac4ad9fcbe354d", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len, 97280, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a", 416, "c7fd0a259771c1c208fb20b3f7a0bb9859a47b6a655e47a4b03ad8a4567f32ee", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len, 97280, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a", 384, "5c3392364b2d53496a1505a6bcb3b179ee718a91aba166d1352a89e6f29447e1", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len, 97280, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a", 416, "8cc703d5b28e8bb936a4ca647f749aba2b49a41cbc8c33dcc874b54b074a0f8b", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len, 97280, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a", 384, "bc9c55c5b547cac8ecceb6a48b3efa2f53c58fd7b1ffb8f8cb0a232c40af29af", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len, 97280, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a", 416, "14d6b39d1ca45a50381c359b6c5cbc432e5b61f4a3e5aa90ed1769eb5d8c394c", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len, 97280, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a", 384, "fa0384182891b60447683a86f986a9d1193388760b3cd27c8440781bccb513ce", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len, 123904, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a", 416, "d687e41aad68fcc7fcb695bf62dcac712024fa879e7fd370cbc55370e542e940", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len, 123904, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a", 384, "eb7cf8dfe1dc909900275dc5ff5693591e5546eaa195941336ee4207dc54aa67", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len, 123904, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a", 416, "a14678fe5f84b7b8541398c5bc5fdb732cadc8d974dd985e6a907cf195c64813", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len, 123904, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a", 384, "12946559dc5d971ad4ded3d1faa415eec7b395ff8b3e55b5f0da8b372de0dd18", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len, 123904, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a", 416, "648ef4107c15ad51db68a834255b4fc43e099b684f2ec5df19d7ba7473e26eae", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len, 123904, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a", 384, "9a90dbaa0accef76e6b3cce8c8a5ad4339bed854a69bfd186519652af3cce24b", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len, 123904, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a", 416, "8da6a223f8df73432b0601dc953b4ecd3542a2e499ba1ccd36a9dabb8d5774f2", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len, 123904, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a", 384, "5dbb3a228c68dfc4cef3a5c987306fa6e387b801f133d57b897bce0626fb07cc", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len, 123904, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a", 416, "ec38064002c223d2f08fa823db4f902e671f5c1b7029e70ff786a0a08c129341", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len, 123904, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a", 384, "dceea69d8a13feddcb5a63e3300afd3760bbd4186de64294f0452243fae2a8b7", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len, 123904, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a", 416, "7cc2536574de984495b44de0ef8f6c45f6381f9abdaab01f0cc367c0a8661a69", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len, 123904, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a", 384, "0a74ae88978f2aad22d47816aeb36e1229b398cddb63d2a1c00266e33f98ebb3", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len, 123904, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a", 416, "dd18f671d972032a78dd12641870f989cb3b8df3176fd867331224f492d74227", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len, 123904, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a", 384, "597374f4d8c6227a98a134bb20ff6e45491f6f994d44633b4d8a339e8c50cdb9", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len, 123904, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a", 416, "cb8335a75a566e201ce08b7bd12a270adaa71f0c0a7ced206684fb8c544f891b", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len, 123904, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a", 384, "c2ad2912f7315a2abfddc2e8a2f007306ff689da1475a60b7b8e219a97dd2fbe", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_sm100a_cubin_len, 61440, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_sm100a", 384, "027475aef4ed9229b874e3b01dfd136f909f1c1d0bb3fc382e283f944eb7fe06", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 1 +, /* mNumTokens */ 0 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a_cubin_len, 61440, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a", 384, "98860c3803b347f64ce0a07d13e1e713630b45a01c705a70eee3611ba0e650b7", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 1 +, /* mNumTokens */ 0 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_sm100a_cubin_len, 61440, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_sm100a", 384, "4119c179a085bacf3f07d57cd888269cb12b8cb0e8ea12d50d1add78ffdfbb4c", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 1 +, /* mNumTokens */ 0 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a_cubin_len, 61440, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a", 384, "6475afd431e333182fd56f19060a99c16dd676c15115a99b877ca582737e269e", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 1 +, /* mNumTokens */ 0 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len, 149504, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a", 416, "d27479f40eb777039050cf3452590e2a0c1fd35437b1535fba4861b89874e677", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 8 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len, 149504, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a", 384, "28c787a8fde1474cb1c432c3299d42896b1660cc987b1696671819ed2ce96b2b", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 8 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len, 149504, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a", 416, "47dde169efefa34d2faba0364d087d1fdf986134d9e1de0764aa733354336d2d", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 8 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len, 149504, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a", 384, "e53508cca19d5b67ad3cc789f344377bb8896494c05bfcd8bb90c01198b71899", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 8 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len, 149504, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a", 416, "5fb79245cf1df99a0cf864e665c2098da795adbd377e2e4bdd24756e30661744", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 8 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len, 149504, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a", 384, "69abbc9a29f37e190aed3fcc1c7a2cb4f2c0ced73b42a54e7f3b8d84d0ec508a", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 8 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len, 149504, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a", 416, "3ce175c6dfbe2ed1296a15a5573169ef8f3630e81191fa8bd1183784a6907ada", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 8 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len, 149504, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a", 384, "43bfe6367d451c1641749b52244ae40ed9bc084d9ce825e78705e42cee76f7c9", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 8 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_sm100a_cubin_len, 61440, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_sm100a", 384, "d7c94aa9865646a2d3d7b36b7daa9b06b6145688924532f6f1fcd651ea1e448e", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 1 +, /* mNumTokens */ 0 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a_cubin_len, 61440, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a", 384, "27f5bed11c1be2270461d7a16bd04c2d5b0c3d1da365926a02b694749f047fbf", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 1 +, /* mNumTokens */ 0 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_sm100a_cubin_len, 61440, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_sm100a", 384, "ef5d746adc2a09e3294b52ce6684d33f8cef0fc6734fe70e38efdfc105642e5c", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 1 +, /* mNumTokens */ 0 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a_cubin_len, 61440, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a", 384, "10a957da353c29e2cd57852d7dcd90d3fea0980ce130d3ef3d5ac9f6d27f742c", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 1 +, /* mNumTokens */ 0 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len, 149504, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a", 416, "f8c169eddeb2274d509a10fcfed8754a2b9657267d56e804da0dcc5f72ba54a2", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 8 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len, 149504, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a", 384, "6fe46321db31e54185cbc03e20e154fb7895c0ad1888365b9d51cedd63d5764c", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 8 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len, 149504, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a", 416, "3ff9f922cf74e03cb5b3f03172d4be34ceadf911ff3b9ce6d6b646ace5eb9a77", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 8 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len, 149504, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a", 384, "b4c54f4f30f019d312346106faa8384065f7e99e9803e46343e7838a1a818664", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 8 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len, 149504, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_clmp_dynBatch_sm100a", 416, "790b2b58ace3c7a9365e8f37f34de8e2b6dabb5c7742cdb33c06d16484383ae9", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 8 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len, 149504, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_dynBatch_sm100a", 384, "a916d59461e1c44f8897a40d017f42b68f1fa86e93679517fef80c1bc49f043d", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 8 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a_cubin_len, 149504, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_clmp_dynBatch_sm100a", 416, "e40531b337de3dff738fd228d732675485ba0e32a3a0f4cfba34f106edce258a", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 8 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a_cubin_len, 149504, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_dynBatch_sm100a", 384, "53d992332e02a59947bd40929f9c8efd3da79234a77d6fddabd868d58bdff5d5", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 8 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin_len, 217088, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a", 256, "b9f33b3ab29d19f2db6159a55ac761b1e4b128b06f6d7108881aeff8ae59dfff", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin_len, 217088, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a", 224, "1db2412318ecc631ab85bd07feaf9564167ccfa09f89d4aa61110f221fb6e279", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a_cubin_len, 217088, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a", 224, "9c7d64b40f4417659ebc3447b30a198f2eeff2bad4e30453f0dc3a00656d33db", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 1 +, /* mNumTokens */ 0 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a_cubin_len, 217088, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_bN_clmp_dynBatch_sm100a", 256, "6f2017653064416fc97c0dda5e5070d3d72e07e2a7c423f7ff085162898f08ed", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a_cubin_len, 217088, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_dynBatch_sm100a", 224, "7313e5403903ff8374ebe03bc3069a248caeab61f28ae356afcde042a4b13996", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a_cubin, Bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a_cubin_len, 217088, "bmm_Bfloat16_E4m3E4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a", 224, "16180b696decf4306a25573ae19dad86217014aecd96aaa44de1b513f3acb422", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 1 +, /* mNumTokens */ 0 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x16x256_s3_et128x16_m128x16x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x16x256_s3_et128x16_m128x16x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 88064, "bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x16x256_s3_et128x16_m128x16x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a", 416, "52b258085b7e0970dc5268ec7ca7ca3a5c4980b94e2b63991fae61dac9f2a763", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1052672) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1052672) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 16 +, /* mMmaKind */ trtllm::gen::MmaKind(1) +, /* mMmaM */ 128 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 1 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x16x256_s3_et128x16_m128x16x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x16x256_s3_et128x16_m128x16x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 88064, "bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x16x256_s3_et128x16_m128x16x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 448, "fe0e4a968b4472f8856c2198f59dfb3fe778e733b166fc0fcc042d822964ceac", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1052672) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1052672) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 16 +, /* mMmaKind */ trtllm::gen::MmaKind(1) +, /* mMmaM */ 128 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 1 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x16x256_s3_et128x16_m128x16x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x16x256_s3_et128x16_m128x16x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 88064, "bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x16x256_s3_et128x16_m128x16x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a", 384, "3b112df01a3b928b1b1bb19f278011a18e9b866715cd19887ccbe200dac218d9", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1052672) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1052672) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 16 +, /* mMmaKind */ trtllm::gen::MmaKind(1) +, /* mMmaM */ 128 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 1 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x16x256_s3_et128x16_m128x16x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x16x256_s3_et128x16_m128x16x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 88064, "bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x16x256_s3_et128x16_m128x16x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 416, "d4aa941092162be422042c26e6868843f229c8d3f374a6760d0627922d0105c0", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1052672) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1052672) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 16 +, /* mMmaKind */ trtllm::gen::MmaKind(1) +, /* mMmaM */ 128 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 1 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x16x256u2_s3_et128x16_m128x16x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x16x256u2_s3_et128x16_m128x16x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 88064, "bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x16x256u2_s3_et128x16_m128x16x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a", 416, "f3edac4357dc88e43c54167c4314f321943c0e63eea05c8234ee121074b57e3f", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1052672) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1052672) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 16 +, /* mMmaKind */ trtllm::gen::MmaKind(1) +, /* mMmaM */ 128 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 1 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x16x256u2_s3_et128x16_m128x16x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x16x256u2_s3_et128x16_m128x16x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 88064, "bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x16x256u2_s3_et128x16_m128x16x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 448, "fa31b14ddda97d48a48d09d14d3c9ce86fdf9285344875bb3a8ef7a717c8bc9a", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1052672) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1052672) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 16 +, /* mMmaKind */ trtllm::gen::MmaKind(1) +, /* mMmaM */ 128 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 1 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x16x256u2_s3_et128x16_m128x16x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x16x256u2_s3_et128x16_m128x16x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 88064, "bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x16x256u2_s3_et128x16_m128x16x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a", 384, "dfeda97aa906fed9d8564bf3c1ddfbde59633837fb730f60006fb593ac60d51b", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1052672) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1052672) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 16 +, /* mMmaKind */ trtllm::gen::MmaKind(1) +, /* mMmaM */ 128 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 1 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x16x256u2_s3_et128x16_m128x16x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x16x256u2_s3_et128x16_m128x16x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 88064, "bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x16x256u2_s3_et128x16_m128x16x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 416, "cc440cac9aacdd9008b8e8b1ab1040459eda7702fdbd2f895bb4bc926bd2fc4f", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1052672) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1052672) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 16 +, /* mMmaKind */ trtllm::gen::MmaKind(1) +, /* mMmaM */ 128 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 1 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x32x256_s3_et128x32_m128x32x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x32x256_s3_et128x32_m128x32x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 116736, "bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x32x256_s3_et128x32_m128x32x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a", 416, "06048dcfaaa81ff67c1f72634fdfbbafb81092dfef6bbb28c7b75e1d57b4ff2a", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1052672) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1052672) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 16 +, /* mMmaKind */ trtllm::gen::MmaKind(1) +, /* mMmaM */ 128 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 1 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x32x256_s3_et128x32_m128x32x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x32x256_s3_et128x32_m128x32x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 116736, "bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x32x256_s3_et128x32_m128x32x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 448, "b85032d1fc6b97f44cd9d3778cba54b25bdcf1cedd90efe34d15f8ecf6eb691e", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1052672) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1052672) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 16 +, /* mMmaKind */ trtllm::gen::MmaKind(1) +, /* mMmaM */ 128 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 1 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x32x256_s3_et128x32_m128x32x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x32x256_s3_et128x32_m128x32x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 116736, "bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x32x256_s3_et128x32_m128x32x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a", 384, "9a52732e9b7fc056e74775739c71b9659f9561386dcfb7ff18a13f634ca59ced", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1052672) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1052672) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 16 +, /* mMmaKind */ trtllm::gen::MmaKind(1) +, /* mMmaM */ 128 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 1 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x32x256_s3_et128x32_m128x32x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x32x256_s3_et128x32_m128x32x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 116736, "bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x32x256_s3_et128x32_m128x32x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 416, "448185b1fd23f3315d5b1fd2ef833066d53184700117b86fa3b6bfa33a820c5c", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1052672) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1052672) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 16 +, /* mMmaKind */ trtllm::gen::MmaKind(1) +, /* mMmaM */ 128 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 1 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x32x256u2_s3_et128x32_m128x32x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x32x256u2_s3_et128x32_m128x32x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 116736, "bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x32x256u2_s3_et128x32_m128x32x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a", 416, "84ea4dd3bb574e7668780351d3c33c7846921eb7fc2a00794a7cbc91c96f7ba4", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1052672) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1052672) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 16 +, /* mMmaKind */ trtllm::gen::MmaKind(1) +, /* mMmaM */ 128 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 1 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x32x256u2_s3_et128x32_m128x32x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x32x256u2_s3_et128x32_m128x32x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 116736, "bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x32x256u2_s3_et128x32_m128x32x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 448, "4cde38e9d6178a2da1028c81bb7897fd0ed5be177cd949ce7f8887ef1f0bd811", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1052672) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1052672) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 16 +, /* mMmaKind */ trtllm::gen::MmaKind(1) +, /* mMmaM */ 128 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 1 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x32x256u2_s3_et128x32_m128x32x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x32x256u2_s3_et128x32_m128x32x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 116736, "bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x32x256u2_s3_et128x32_m128x32x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a", 384, "79bbbc653e8be33b7e8b2ca2400a7e2a9eb94752912f18fc291a0d3b6407e979", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1052672) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1052672) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 16 +, /* mMmaKind */ trtllm::gen::MmaKind(1) +, /* mMmaM */ 128 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 1 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x32x256u2_s3_et128x32_m128x32x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x32x256u2_s3_et128x32_m128x32x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 116736, "bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x32x256u2_s3_et128x32_m128x32x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 416, "5c1d3679cb47511def1d4db96988ba8104914592a46b4ad3df554e2ffa482577", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1052672) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1052672) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 16 +, /* mMmaKind */ trtllm::gen::MmaKind(1) +, /* mMmaM */ 128 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 1 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x64x256_s3_et128x64_m128x64x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x64x256_s3_et128x64_m128x64x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 175104, "bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x64x256_s3_et128x64_m128x64x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a", 416, "5c579bd1d00f0f0ab698aa5cfd70596283f367c962081c707b7b8f101ab84793", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1052672) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1052672) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 16 +, /* mMmaKind */ trtllm::gen::MmaKind(1) +, /* mMmaM */ 128 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 1 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x64x256_s3_et128x64_m128x64x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x64x256_s3_et128x64_m128x64x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 175104, "bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x64x256_s3_et128x64_m128x64x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 448, "d2800340cfeda9ccccbcc4951b64233263b58bdc12e79896f61ea26cb41edbab", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1052672) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1052672) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 16 +, /* mMmaKind */ trtllm::gen::MmaKind(1) +, /* mMmaM */ 128 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 1 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x64x256_s3_et128x64_m128x64x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x64x256_s3_et128x64_m128x64x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 175104, "bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x64x256_s3_et128x64_m128x64x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a", 384, "469c956f7e2e095908ddd733b1db48d1effbff506bca5efc43f7395d72c1c48f", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1052672) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1052672) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 16 +, /* mMmaKind */ trtllm::gen::MmaKind(1) +, /* mMmaM */ 128 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 1 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x64x256_s3_et128x64_m128x64x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x64x256_s3_et128x64_m128x64x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 175104, "bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x64x256_s3_et128x64_m128x64x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 416, "13efa9ae5dd6d8e88d98c2f5af777eef65a538a3530cccfa08054c6e26714c93", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1052672) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1052672) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 16 +, /* mMmaKind */ trtllm::gen::MmaKind(1) +, /* mMmaM */ 128 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 1 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x64x256u2_s3_et128x64_m128x64x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x64x256u2_s3_et128x64_m128x64x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 175104, "bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x64x256u2_s3_et128x64_m128x64x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a", 416, "a6e68338e9fe01626be6c254deda824bf588a69be035e195254de082ca68c723", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1052672) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1052672) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 16 +, /* mMmaKind */ trtllm::gen::MmaKind(1) +, /* mMmaM */ 128 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 1 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x64x256u2_s3_et128x64_m128x64x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x64x256u2_s3_et128x64_m128x64x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 175104, "bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x64x256u2_s3_et128x64_m128x64x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 448, "28fae92ddb56b8569597955cef7b840c995d6d8c375bd0b6404086dfe722b025", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1052672) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1052672) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 16 +, /* mMmaKind */ trtllm::gen::MmaKind(1) +, /* mMmaM */ 128 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 1 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x64x256u2_s3_et128x64_m128x64x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x64x256u2_s3_et128x64_m128x64x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 175104, "bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x64x256u2_s3_et128x64_m128x64x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a", 384, "c4a5d32e52bc38ebc69fd0102bd1c4f4444f114345b24674832acae0e6b34a11", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1052672) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1052672) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 16 +, /* mMmaKind */ trtllm::gen::MmaKind(1) +, /* mMmaM */ 128 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 1 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x64x256u2_s3_et128x64_m128x64x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x64x256u2_s3_et128x64_m128x64x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 175104, "bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x64x256u2_s3_et128x64_m128x64x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 416, "9763f6a7f20c8fa74157036a62a755dc6569ebee9cf054d1c088182633bb234e", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1052672) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1052672) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 16 +, /* mMmaKind */ trtllm::gen::MmaKind(1) +, /* mMmaM */ 128 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 1 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x8x256_s3_et128x8_m128x8x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x8x256_s3_et128x8_m128x8x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 73728, "bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x8x256_s3_et128x8_m128x8x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a", 416, "3a28fe8644a7bc7a1cc014a8897ab0fe2166436e2cd5cfaed18d8fb57ad69413", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1052672) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1052672) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 16 +, /* mMmaKind */ trtllm::gen::MmaKind(1) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 1 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x8x256_s3_et128x8_m128x8x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x8x256_s3_et128x8_m128x8x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 73728, "bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x8x256_s3_et128x8_m128x8x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 448, "fae33ca8316876b52ad4442738631714bb966dce0092b74026b4b1186ec074b3", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1052672) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1052672) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 16 +, /* mMmaKind */ trtllm::gen::MmaKind(1) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 1 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x8x256_s3_et128x8_m128x8x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x8x256_s3_et128x8_m128x8x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 73728, "bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x8x256_s3_et128x8_m128x8x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a", 384, "84acc24ca690945c29621f601d016a8dd6a0d1469b8d735806fc726b41bd1793", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1052672) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1052672) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 16 +, /* mMmaKind */ trtllm::gen::MmaKind(1) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 1 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x8x256_s3_et128x8_m128x8x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x8x256_s3_et128x8_m128x8x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 73728, "bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x8x256_s3_et128x8_m128x8x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 416, "b00a33a971a83513b638849fb432d0fbe11eb4b38ae4a696227295964b939aa5", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1052672) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1052672) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 16 +, /* mMmaKind */ trtllm::gen::MmaKind(1) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 1 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x8x256u2_s3_et128x8_m128x8x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x8x256u2_s3_et128x8_m128x8x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 73728, "bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x8x256u2_s3_et128x8_m128x8x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a", 416, "9e7147c1111d5312284b3711b5cb5e5615212529b87d67cf319d923276dc7d31", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1052672) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1052672) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 16 +, /* mMmaKind */ trtllm::gen::MmaKind(1) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 1 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x8x256u2_s3_et128x8_m128x8x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x8x256u2_s3_et128x8_m128x8x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 73728, "bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x8x256u2_s3_et128x8_m128x8x16_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 448, "15abd2a84edab7a22acf98b74499a0789fd767a68a813e147f6ecd3372f0f851", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1052672) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1052672) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 16 +, /* mMmaKind */ trtllm::gen::MmaKind(1) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 1 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x8x256u2_s3_et128x8_m128x8x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x8x256u2_s3_et128x8_m128x8x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 73728, "bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x8x256u2_s3_et128x8_m128x8x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a", 384, "b386b8c7e1c9339659e6837240fa2f21f8fe927dd5374fce061d57517c9fcf1e", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1052672) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1052672) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 16 +, /* mMmaKind */ trtllm::gen::MmaKind(1) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 1 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x8x256u2_s3_et128x8_m128x8x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x8x256u2_s3_et128x8_m128x8x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 73728, "bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_patch_Fp32_t128x8x256u2_s3_et128x8_m128x8x16_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 416, "ef2e3792a633ceba03139e91beb01517ee09d2b7539f67e146be12104e284280", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1052672) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1052672) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 16 +, /* mMmaKind */ trtllm::gen::MmaKind(1) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 1 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x16x256_s5_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x16x256_s5_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 201728, "bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x16x256_s5_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a", 320, "38b5f0e760783280d227c998b300f0e8dd71a3f26373714a22e2f31348a3e5ed", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 5 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x16x256_s5_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x16x256_s5_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 201728, "bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x16x256_s5_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a", 288, "a3a12a8fad98acfb1dee2a7c4d869a39890ffce9d5db5966da9a82a46da2ed62", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 5 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x16x256u2_s5_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x16x256u2_s5_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 201728, "bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x16x256u2_s5_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a", 320, "6b3293be7d956ba7b71371bea3abaaedd83b8f579db9a059c6fc99693223c4cd", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 5 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x16x256u2_s5_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x16x256u2_s5_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 201728, "bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x16x256u2_s5_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a", 288, "22b5ab9b379205e9fc715a447023dbd202303b0f8d491ed3bffc77c33be2587c", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 5 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x32x256_s5_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x32x256_s5_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 226304, "bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x32x256_s5_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a", 320, "ef6fd532b1f40ca268e0ae2f2dd2c862ebf43e67bf57058d889346a0e25f1725", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 5 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x32x256_s5_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x32x256_s5_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 226304, "bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x32x256_s5_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a", 288, "33a419aa9f7d3efe0f05b76483b1f745471a65cf59cfb76946d3abca189152af", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 5 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x32x256u2_s5_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x32x256u2_s5_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 226304, "bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x32x256u2_s5_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a", 320, "169635b45a94a5ac3d7abc68ac2d3de3cd5438afab71b672b4f2eec22f16cdd0", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 5 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x32x256u2_s5_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x32x256u2_s5_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 226304, "bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x32x256u2_s5_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a", 288, "839da645ae7ddebfa72c9a73ee25bc0af44c233f92744a74ae2b3a0853caceb6", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 5 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x64x256_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x64x256_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 225280, "bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x64x256_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a", 320, "0994bfa66ea8a078b03947f88850366fca52db246742509924bf7a6f213c7bfb", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x64x256_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x64x256_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 225280, "bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x64x256_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a", 288, "ec379b195e72575bf96f9ff16ea0c35eb91e9b291a5a3719751407d6daed6337", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x64x256u2_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x64x256u2_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 225280, "bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x64x256u2_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a", 320, "5b5d1151de0658df84e2810b3f1916d297be4ce3651d95d9a9de889e4afd1c86", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x64x256u2_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x64x256u2_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 225280, "bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x64x256u2_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a", 288, "3b0858a36167ad32d9905500e2aa59743b6d91bb2c2373a70f9c3814aa074795", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 189440, "bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a", 320, "91874a966f6eb543323f8467c2e3f66600f20f0ab8754026abde6b510c699f98", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 5 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 189440, "bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a", 288, "955ae7dd9ea62cf762ca07c8ee875e24303b0d95b46b03768c07ad3631e5b143", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 5 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 189440, "bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a", 320, "e11a6189ded24497219f2bf335f60a83fb17e163dc32459c7d166a015e2720be", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 5 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 189440, "bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a", 288, "be1acac35d67130541d30ead271c206f7ac6f70ac4a0e46db4c1845053a8c3a1", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 5 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 225280, "bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a", 320, "0033ff8a409809884ee01fd3837975710d5cf654422b8055892c126cfe11419d", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 225280, "bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a", 288, "7ad06f231f869ea9e835d3d1083e027d7ecca6f7a5fb03752f9d351f7ae2fc45", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 225280, "bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a", 320, "620b6d07cfbcf2cfbbc16d4e4dedaba7eccdac2f878dda447431a34594cb95d4", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 225280, "bmm_Bfloat16_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a", 288, "4e8fea6373057cb4cee2f2a32a4d0edee7817c060b3449471b0001bc9dd2d482", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x16x256_s3_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x16x256_s3_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 125952, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x16x256_s3_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a", 480, "0b7ed55068401563c1d4f1b8d0decb49821a0aaa2151c02ce4feca7ca038329b", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x16x256_s3_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x16x256_s3_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 125952, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x16x256_s3_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a", 448, "198533be3cee3890a08a67e379a1223f4daec31058c520726cdd29bdaf13560a", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x16x256_s4_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x16x256_s4_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 163840, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x16x256_s4_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a", 480, "f1fd3a527c984a2775575ada2a9f8c3dc5b1a67a16e2f8c33b7a5536cdc3a00e", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x16x256_s4_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x16x256_s4_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 163840, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x16x256_s4_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a", 448, "d54f4f208b17054371d86be133399e8eca22070d2698208cb8706b537baa67a4", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x16x256u2_s3_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x16x256u2_s3_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 125952, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x16x256u2_s3_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a", 480, "be7f908008125834306e4050d77214cb231f8f405e7eee09d01defa33a3cf2f8", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x16x256u2_s3_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x16x256u2_s3_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 125952, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x16x256u2_s3_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a", 448, "9a71d9dabf06f8c485511ffbf2596ba8c56d05acc22f82fc16b3f53dc35d74a7", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x16x256u2_s4_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x16x256u2_s4_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 163840, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x16x256u2_s4_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a", 480, "ab7b410bd93cefea6d39fb831410c13fa2b9f72d772b7055f9a504f4e0b9cef9", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x16x256u2_s4_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x16x256u2_s4_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 163840, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x16x256u2_s4_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a", 448, "bb56c004e9600811b1287fd79ee6b2d5ba1a6111bbebb3a01590b4a52042dcb7", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x32x256_s3_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x32x256_s3_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 142336, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x32x256_s3_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a", 480, "3dd7d632ae2bd741bd66105cca2ff630e928956a57cb269180ace28ed0a2bb9a", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x32x256_s3_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x32x256_s3_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 142336, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x32x256_s3_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a", 448, "675e511eb3d5efc57b05609cfa044a3228077d83033845d272790997b613346f", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x32x256_s4_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x32x256_s4_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 185344, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x32x256_s4_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a", 480, "13f1e99d4a76143862bac3e5f420a10587bd6be5286615de4aad9b2ead3948c7", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x32x256_s4_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x32x256_s4_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 185344, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x32x256_s4_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a", 448, "9bbfd28c75e78ae54f4d2b58aff19bea505e96dc397d3dfc5031c815601826ea", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x32x256u2_s3_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x32x256u2_s3_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 142336, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x32x256u2_s3_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a", 480, "712986701528a94354455e39c9759c5573c463a8a47df822c2286a65fe0daefb", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x32x256u2_s3_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x32x256u2_s3_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 142336, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x32x256u2_s3_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a", 448, "ea38d7b3c60b20595e79c692903917e5cee01bfc2ac9e35af00ca03286e62956", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x32x256u2_s4_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x32x256u2_s4_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 185344, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x32x256u2_s4_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a", 480, "a1c9af9aa0f57eac713088b8b31f5b216a145c503cdf7944880bebdde359dbe8", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x32x256u2_s4_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x32x256u2_s4_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 185344, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x32x256u2_s4_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a", 448, "aef9a91ee0738e331f8d7c3bb3faa862113a247b9754ba6e12ac9766d452eb27", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x64x256_s3_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x64x256_s3_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 177152, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x64x256_s3_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a", 480, "64f070388a1a75610038a835e814097a80e468aa6386be16006bc24289a245e9", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x64x256_s3_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x64x256_s3_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 177152, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x64x256_s3_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a", 448, "1b6974f8108e1b0cdc6e3b066b84e0ec4cd580cfd3267ac1ec941bda41d954c8", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x64x256_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x64x256_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 228352, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x64x256_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a", 480, "553f0aa838c31e24f41ac380d797de027c41358352544642cf585919e69b22dc", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x64x256_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x64x256_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 228352, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x64x256_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a", 448, "5336f9f353d3b5d2a204a95b22ea52c774bd6362ea2659a53a401062512cea56", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x64x256u2_s3_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x64x256u2_s3_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 177152, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x64x256u2_s3_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a", 480, "587c54db90f77353dc7eaaf0130bc816e1acdc27ff84207b7ae45af5f203e63f", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x64x256u2_s3_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x64x256u2_s3_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 177152, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x64x256u2_s3_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a", 448, "79c219a8c1326757393006cdaf72de6389a5aa9a3dda955141b54e4f61e832b4", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x64x256u2_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x64x256u2_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 228352, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x64x256u2_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a", 480, "e2e741c1d17440c4463d8f941baa4dbd9d63ac2ef5f5183df5ad1bf25cf06491", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x64x256u2_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x64x256u2_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 228352, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x64x256u2_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a", 448, "3cc64b02c875b4660ee70fa2db4029bee0d7c60eb95446852a6732150133b6e1", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s4_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s4_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 153600, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s4_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a", 480, "1fa7ae2040138a8e919c787e62ff4d5e235ac397ff75c78c56ca1eed1ea2c113", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s4_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s4_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 153600, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s4_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a", 448, "b98f704b2003d04bb36e2a9307a7319e6471e7d4ad7afa2080b657c34ee6d2c8", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s4_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s4_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 159744, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s4_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a", 480, "037db60f5ab7b61a4e80ffaaa61f36fc32dba72bacf09334b15e1da592a53332", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 2 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 2 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(2) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s4_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s4_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 151552, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s4_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a", 448, "bdb9e4f703807f6ce8abc9ae7bc8e4b8bc6f1dc4d026a35d482d0ba0ea2e85aa", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 2 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 2 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(2) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 189440, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a", 480, "e89cf7342fc0013e957c2f58ddb53f84d05d679462c152378ad047134817a689", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 5 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 189440, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a", 448, "27c64ecb23ece3132d8fa547e9d577d20042063f526bcc73d16a10e9e53481e9", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 5 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 195584, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a", 480, "11d7f6c4c4a8d7f382fed2b2dc7f608743064587b19ae1b96d412d9db106a290", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 2 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 2 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 5 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(2) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 187392, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a", 448, "3df6963d871d8fc2915f6e4be2d899ead0f06d4fb3156487ea39e2febda7bff7", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 2 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 2 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 5 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(2) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s6_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s6_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 225280, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s6_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a", 480, "b419076e6f8be20e164ea1ed8d4770b9aba41a20baeb35ba8e0ff72f9c88f5f7", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s6_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s6_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 225280, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s6_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a", 448, "2bef689b3d1f86873e618973d7aef905630574fb54b043dfd3bd1dd6d14d6f23", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s6_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s6_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 231424, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s6_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a", 480, "e425bf76a44f9c244a8b04add431d3dec324e745c3ab58df1890d42c335f8e80", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 2 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 2 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(2) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s6_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s6_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 223232, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256_s6_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a", 448, "ceec79d213aca4da2dd3a230b3f81fa5d8a39890fa81befae3421de09d5316f8", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 2 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 2 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(2) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s4_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s4_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 153600, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s4_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a", 480, "58962324203c255d56409d2fd7bbbe91872efd46168e157a22ac1ac70b7b42b7", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s4_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s4_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 153600, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s4_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a", 448, "ab07787ab1420a630f9cd816a0151aed9b5b3f6270429b42e9255c227676e6b3", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s4_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s4_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 159744, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s4_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a", 480, "c25077e7fc632546ab3389dace009de63e2972d38b9e77e7b980d010b15822c6", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 2 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 2 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(2) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s4_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s4_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 151552, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s4_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a", 448, "534eb114f2c8c1b77de1f33bab90959259d163280e912429b945302d0e06cf28", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 2 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 2 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(2) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 189440, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a", 480, "1862ac649a69e4b532651012a2dff768d45cbca4cf467ccc4041910eb1de1176", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 5 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 189440, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a", 448, "233baae772c1b23100c09b46a79cedb27f13605556f0bc3e60a8839dde6a39e0", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 5 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 195584, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a", 480, "c21ac81f0799a8365f557430e12ce95f448cdb21eab8fe146533078b644c2d3f", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 2 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 2 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 5 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(2) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 187392, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a", 448, "eb55eb91b197919dcff2a184c1037c07363cf1b5a4c68d33ff53d738266af975", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 2 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 2 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 5 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(2) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s6_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s6_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 225280, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s6_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a", 480, "df7d18206ee418bfc9a636f5c326ea35a6c687fb519c3ed4e8ad1fd4b7e224f4", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s6_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s6_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 225280, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s6_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a", 448, "6377876adb967003b602d3ad5e786c703b1ff72331929d490c28eb0f7c863bc6", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s6_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s6_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 231424, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s6_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a", 480, "615bd354244e6679359592eb69ee4bb1fac02720ea6de3d254f5ef4234a5a208", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 2 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 2 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(2) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s6_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s6_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 223232, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x256u2_s6_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a", 448, "a41e0feccf2584546ec060c487d0bc2f0a5b98ae348cf9c593b51c98ec6fd6d5", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 2 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 2 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(2) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 225280, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a", 480, "e464cee5683c77a9b0a8710a79df18c81c30bd1e1ad70cc1110f508283d74a70", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 225280, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a", 448, "1f1b53b8475b8915cb927f17cb914c8f39d3bea4b19cc2dc08a9cb18b13376c4", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 231424, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a", 480, "b6fedcccdb8c1f4a716e93165cd6dbf826b06889abcd2b5548192dccdc5c2464", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 2 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 2 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(2) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 223232, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a", 448, "6f3ad785800894aaca7d8ff8aa211b31800d2e21975bd72751a81f0687e3f2a7", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 2 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 2 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(2) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 225280, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a", 480, "7399ce8bb98825c9f4dc063932e7bde7d4d32aebff3023d53fca90155fb42fef", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 225280, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a", 448, "cd7dd69304a1f7208ba22a4dc2630741bf84a4ec67a664abdb7528a30bce4511", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 231424, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_clmp_dynBatch_sm100a", 480, "3dc4e382e28111d76580bd38dc1af4da201b2dea9f0fd5bb1cc0adf35b9d8d81", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 2 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 2 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(2) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin, Bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a_cubin_len, 223232, "bmm_Bfloat16_MxE2m1MxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_clmp_dynBatch_sm100a", 448, "9048f79c720a2ce32103f677813c48ff741b475431290107f693960b4144637b", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 2 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(1052672) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 2 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(2) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 0 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E2m1_E2m1E2m1_Fp32_t128x16x512_s4_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_E2m1_E2m1E2m1_Fp32_t128x16x512_s4_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 175104, "bmm_E2m1_E2m1E2m1_Fp32_t128x16x512_s4_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 512, "da31f6d9c67d08e85b7072d5a43266d0d6ee30735817b37f982c0477de645a06", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826818) +, /* mDtypeB */ trtllm::gen::Dtype(17826818) +, /* mDtypeC */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17826818) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaKind */ trtllm::gen::MmaKind(4) +, /* mMmaM */ 128 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E2m1_E2m1E2m1_Fp32_t128x16x512_s4_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_E2m1_E2m1E2m1_Fp32_t128x16x512_s4_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 175104, "bmm_E2m1_E2m1E2m1_Fp32_t128x16x512_s4_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 480, "42e9065be76998f1a64683085d9ee72071e85c3f6051fc9aec97b8cb4a8bc7cb", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826818) +, /* mDtypeB */ trtllm::gen::Dtype(17826818) +, /* mDtypeC */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17826818) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaKind */ trtllm::gen::MmaKind(4) +, /* mMmaM */ 128 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E2m1_E2m1E2m1_Fp32_t128x16x512u2_s4_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_E2m1_E2m1E2m1_Fp32_t128x16x512u2_s4_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 175104, "bmm_E2m1_E2m1E2m1_Fp32_t128x16x512u2_s4_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 512, "46c62f5dc36d3a412416fb040ad22f3bbd4d65f5706fd2f32672bb629dc6dd18", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826818) +, /* mDtypeB */ trtllm::gen::Dtype(17826818) +, /* mDtypeC */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17826818) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaKind */ trtllm::gen::MmaKind(4) +, /* mMmaM */ 128 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E2m1_E2m1E2m1_Fp32_t128x16x512u2_s4_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_E2m1_E2m1E2m1_Fp32_t128x16x512u2_s4_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 175104, "bmm_E2m1_E2m1E2m1_Fp32_t128x16x512u2_s4_et128x16_m128x16x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 480, "8e5a000b9df97eed3710350f1126dc698bb7f2cac3ca94e063b469769a176366", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826818) +, /* mDtypeB */ trtllm::gen::Dtype(17826818) +, /* mDtypeC */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17826818) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaKind */ trtllm::gen::MmaKind(4) +, /* mMmaM */ 128 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E2m1_E2m1E2m1_Fp32_t128x32x512_s4_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_E2m1_E2m1E2m1_Fp32_t128x32x512_s4_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 194560, "bmm_E2m1_E2m1E2m1_Fp32_t128x32x512_s4_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 512, "5d621271aba16b84dfe6283b5d93ac500800b567bff705d8a73d2d2a963e67bd", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826818) +, /* mDtypeB */ trtllm::gen::Dtype(17826818) +, /* mDtypeC */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17826818) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaKind */ trtllm::gen::MmaKind(4) +, /* mMmaM */ 128 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E2m1_E2m1E2m1_Fp32_t128x32x512_s4_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_E2m1_E2m1E2m1_Fp32_t128x32x512_s4_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 194560, "bmm_E2m1_E2m1E2m1_Fp32_t128x32x512_s4_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 480, "17abea2a2e8181c85657a1c3654da4d5783bb186fc0547237b6cb88dff148d31", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826818) +, /* mDtypeB */ trtllm::gen::Dtype(17826818) +, /* mDtypeC */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17826818) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaKind */ trtllm::gen::MmaKind(4) +, /* mMmaM */ 128 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E2m1_E2m1E2m1_Fp32_t128x32x512u2_s4_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_E2m1_E2m1E2m1_Fp32_t128x32x512u2_s4_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 194560, "bmm_E2m1_E2m1E2m1_Fp32_t128x32x512u2_s4_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 512, "ef8e4838f68f4020afda48c29d032ea4f56239df6b101d1b8f6c3a15f85d744c", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826818) +, /* mDtypeB */ trtllm::gen::Dtype(17826818) +, /* mDtypeC */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17826818) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaKind */ trtllm::gen::MmaKind(4) +, /* mMmaM */ 128 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E2m1_E2m1E2m1_Fp32_t128x32x512u2_s4_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_E2m1_E2m1E2m1_Fp32_t128x32x512u2_s4_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 194560, "bmm_E2m1_E2m1E2m1_Fp32_t128x32x512u2_s4_et128x32_m128x32x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 480, "f306d4c0b809848b99d22554426480770767cef45047dec3e80f8b161f2c140d", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826818) +, /* mDtypeB */ trtllm::gen::Dtype(17826818) +, /* mDtypeC */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17826818) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaKind */ trtllm::gen::MmaKind(4) +, /* mMmaM */ 128 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E2m1_E2m1E2m1_Fp32_t128x64x256_s6_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_E2m1_E2m1E2m1_Fp32_t128x64x256_s6_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 178176, "bmm_E2m1_E2m1E2m1_Fp32_t128x64x256_s6_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 512, "39dc8af745963ed8342183716ed83d62266c28ddbe2290a00b1401de33537e42", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826818) +, /* mDtypeB */ trtllm::gen::Dtype(17826818) +, /* mDtypeC */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17826818) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaKind */ trtllm::gen::MmaKind(4) +, /* mMmaM */ 128 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E2m1_E2m1E2m1_Fp32_t128x64x256_s6_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_E2m1_E2m1E2m1_Fp32_t128x64x256_s6_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 178176, "bmm_E2m1_E2m1E2m1_Fp32_t128x64x256_s6_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 480, "932c455add6e5927757c2a56c4adfc81ad2dc114bd0c55924950dfc424432fd8", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826818) +, /* mDtypeB */ trtllm::gen::Dtype(17826818) +, /* mDtypeC */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17826818) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaKind */ trtllm::gen::MmaKind(4) +, /* mMmaM */ 128 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E2m1_E2m1E2m1_Fp32_t128x64x256u2_s6_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_E2m1_E2m1E2m1_Fp32_t128x64x256u2_s6_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 178176, "bmm_E2m1_E2m1E2m1_Fp32_t128x64x256u2_s6_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 512, "7492c306444a92b35ff1666b496b0b518b1e5e1a3bf571334d4e29ab6129666f", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826818) +, /* mDtypeB */ trtllm::gen::Dtype(17826818) +, /* mDtypeC */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17826818) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaKind */ trtllm::gen::MmaKind(4) +, /* mMmaM */ 128 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E2m1_E2m1E2m1_Fp32_t128x64x256u2_s6_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_E2m1_E2m1E2m1_Fp32_t128x64x256u2_s6_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 178176, "bmm_E2m1_E2m1E2m1_Fp32_t128x64x256u2_s6_et128x64_m128x64x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 480, "c068236aac9e7192e3c904b839834bdda81a1f8b07d2f69dbe032d0d05374822", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826818) +, /* mDtypeB */ trtllm::gen::Dtype(17826818) +, /* mDtypeC */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17826818) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaKind */ trtllm::gen::MmaKind(4) +, /* mMmaM */ 128 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E2m1_E2m1E2m1_Fp32_t128x8x512_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_E2m1_E2m1E2m1_Fp32_t128x8x512_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 165888, "bmm_E2m1_E2m1E2m1_Fp32_t128x8x512_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 512, "e470d3aabe53bc2484d130ce968c0ebb8776d1a2c563377c0e6a128ec42efe2b", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826818) +, /* mDtypeB */ trtllm::gen::Dtype(17826818) +, /* mDtypeC */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17826818) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaKind */ trtllm::gen::MmaKind(4) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E2m1_E2m1E2m1_Fp32_t128x8x512_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a_cubin, Bmm_E2m1_E2m1E2m1_Fp32_t128x8x512_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a_cubin_len, 165888, "bmm_E2m1_E2m1E2m1_Fp32_t128x8x512_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a", 448, "0e1aa67fd306295f4c21d072e436bcdc1d51fef87aa2fa0b0493b91aefc481c4", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826818) +, /* mDtypeB */ trtllm::gen::Dtype(17826818) +, /* mDtypeC */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17826818) +, /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaKind */ trtllm::gen::MmaKind(4) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 1 +, /* mNumTokens */ 0 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E2m1_E2m1E2m1_Fp32_t128x8x512_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_E2m1_E2m1E2m1_Fp32_t128x8x512_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 165888, "bmm_E2m1_E2m1E2m1_Fp32_t128x8x512_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 480, "142afcee25cb84bf6611830b8a71e127deb69232f580a05df56e6ee505a09f05", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826818) +, /* mDtypeB */ trtllm::gen::Dtype(17826818) +, /* mDtypeC */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17826818) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaKind */ trtllm::gen::MmaKind(4) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E2m1_E2m1E2m1_Fp32_t128x8x512u2_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_E2m1_E2m1E2m1_Fp32_t128x8x512u2_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 165888, "bmm_E2m1_E2m1E2m1_Fp32_t128x8x512u2_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedP_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 512, "238320111ff22f9aedabaa9d79954e7db0a395df3c3ff50af7306616307219b1", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826818) +, /* mDtypeB */ trtllm::gen::Dtype(17826818) +, /* mDtypeC */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17826818) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaKind */ trtllm::gen::MmaKind(4) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E2m1_E2m1E2m1_Fp32_t128x8x512u2_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a_cubin, Bmm_E2m1_E2m1E2m1_Fp32_t128x8x512u2_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a_cubin_len, 165888, "bmm_E2m1_E2m1E2m1_Fp32_t128x8x512u2_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a", 448, "ba95baba14a5cb89863871c716d247bef2b461df8c88a1fd6668e782254bb8ee", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826818) +, /* mDtypeB */ trtllm::gen::Dtype(17826818) +, /* mDtypeC */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17826818) +, /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaKind */ trtllm::gen::MmaKind(4) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 1 +, /* mNumTokens */ 0 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E2m1_E2m1E2m1_Fp32_t128x8x512u2_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_E2m1_E2m1E2m1_Fp32_t128x8x512u2_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 165888, "bmm_E2m1_E2m1E2m1_Fp32_t128x8x512u2_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 480, "9ae64768181f24c6ca90e899261eabbf76d93319ca54c212cac0764e5f9f89b7", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826818) +, /* mDtypeB */ trtllm::gen::Dtype(17826818) +, /* mDtypeC */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17826818) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaKind */ trtllm::gen::MmaKind(4) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 119808, "bmm_E4m3_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a", 448, "8157f44703c59388fb09afac934d18fb37a94aec57207801f07c8603687ecdbe", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 119808, "bmm_E4m3_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a", 416, "7477fff17f203078d1a94b4ec28bc0b061c1a1df194da9dfa3e313820f3c010e", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 119808, "bmm_E4m3_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a", 448, "2f8ad526b885e8baa349e382a4ae0ac982d886242594141ae02519979b71b13b", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 119808, "bmm_E4m3_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a", 416, "3e8760064685db77751a0651e21f944c694263647635671dd3958ed3f7fada28", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 119808, "bmm_E4m3_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a", 448, "508b4b08f46d803db6ece157302800429d3a315b448f82c594f09dc224054f50", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 119808, "bmm_E4m3_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a", 416, "fb124fa8b6e12b989155fab1c39e1e9c02f8aaa7b10a48a990e1e69f87baa9d0", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 119808, "bmm_E4m3_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a", 448, "7cec49ee4eb48238caac11db363f407c6af4b5155fc1f422a89ca35ff92d2d7a", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 119808, "bmm_E4m3_E4m3E4m3_Fp32_t128x16x128_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a", 416, "b63a1328798a2b363926f1e908e874ef1a94f5dfa4bdff87f24731d7c5a68636", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 119808, "bmm_E4m3_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a", 448, "bfd97bbe6574a722b831054f05d8f9a60f921836e77fe307bcae58e37198be70", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 119808, "bmm_E4m3_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a", 416, "ab4005174cdc2b54fb322c32c184952ecbcbdfbd99d173d7e242d382da88fdf6", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 119808, "bmm_E4m3_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a", 448, "3c46c6b249e5a0cd7340f3dbc7b8072077652ec4c00f825b388f8b8bbea69cbf", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 119808, "bmm_E4m3_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a", 416, "d38baf93035e39c7f6d3928ac5f5e95037202d5f5a668c47559e47725049baa7", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 119808, "bmm_E4m3_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a", 448, "3d88d7cdf8878e120633daa4940007ff6fd4ea0c78dca63819117adeea74afde", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 119808, "bmm_E4m3_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a", 416, "5fc17978448b4e1bbac488501ae56fbbcdade025754f46e2cb05e94a63b67a82", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 119808, "bmm_E4m3_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a", 448, "2eaeea22b7f06a714c66e5e7cb07a926c9813c6442c7f81b0d21624d9c8a5920", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 119808, "bmm_E4m3_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a", 416, "028d9e988fc094d7ffe0d9b9416236f89fba7f1bcb8a5fbcfb5dc499b2174955", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x16x512_s3_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_tokSfB_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x16x512_s3_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_tokSfB_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 229376, "bmm_E4m3_E4m3E4m3_Fp32_t128x16x512_s3_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_tokSfB_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 256, "f29f44b4a5e9666850a7ae0503f150a76fa4af45d18242fa53652c7830a492d6", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 128 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 1 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x16x512u2_s3_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_tokSfB_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x16x512u2_s3_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_tokSfB_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 229376, "bmm_E4m3_E4m3E4m3_Fp32_t128x16x512u2_s3_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_tokSfB_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 256, "e3c30edacd1a1361c6e983c5724253f6f0963c2af08a722f855e0a476fd02d41", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 128 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 1 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 93184, "bmm_E4m3_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a", 448, "3156acfad50fdc17cad8a8361003459321a3464ced26e0028ac7b0e323b134b6", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 93184, "bmm_E4m3_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a", 416, "96fcdb017b1d06c94709523848b757776a2c00ef601e2b7583f0f6d84e762741", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 93184, "bmm_E4m3_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a", 448, "020a6868ec1616be5d6672af219d2f37babd931ed22ea7f6b351c51be22f2b2c", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 93184, "bmm_E4m3_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a", 416, "d1f59e9d0d394769ff8709dccaf513fdf564efd007962aac9ab9f6cc4ab5a0e7", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 93184, "bmm_E4m3_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a", 448, "7612966658dff79d86fd595ddec5ffe687dad3df5381593a88f837720c236964", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 93184, "bmm_E4m3_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a", 416, "1c6e04f534cadf6f08823442154ac262c67ffcf39f9ccac08a21a4d9d2a901dd", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 93184, "bmm_E4m3_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a", 448, "1133ac386e288984301755cb9ea35960ce9356eff891c4f0e125a3216ea36c3e", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 93184, "bmm_E4m3_E4m3E4m3_Fp32_t128x32x128_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a", 416, "04f3fe5b18e41354a51d94d394f030f5196619ac176e2332c5b58f16fc4ef8c0", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 93184, "bmm_E4m3_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a", 448, "e9a6951414db31e6fe9792c25afee459948dd0d2866a5fc78f6faf0441a52dd1", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 93184, "bmm_E4m3_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a", 416, "be24ef39cdd4e0c681b6e30e5558549487f2db3006a396cbf1a4a71446f99b31", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 93184, "bmm_E4m3_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a", 448, "62c5de8afbf9ae1890e508e6ff4ebbef8db86ad0f25233bfb24e1e894da6860d", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 93184, "bmm_E4m3_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a", 416, "543cd01fcc2f8a4bd4250bf645085f3d3663c50dcc17e1529c13794677b5accf", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 93184, "bmm_E4m3_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a", 448, "6032dad65ba48e4b92f9a2fcaaf48f375f9bdd235b8948be74df9e1befb92ecb", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 93184, "bmm_E4m3_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a", 416, "2a620944a3085c1054f57d0f07c6d19ac9d7958858607658a2961fd0c0eef409", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 93184, "bmm_E4m3_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a", 448, "f7bdd240292559badd385d400b123e051f2487e3e4dc16baf4b74b71cb4b1cef", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 93184, "bmm_E4m3_E4m3E4m3_Fp32_t128x32x128u2_s4_et64x32_m64x32x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a", 416, "72ece78605b6d5fafc6bfbbaaef349293f31d58fe16f1ae1bd548f0754b57c4e", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x32x512_s2_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_tokSfB_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x32x512_s2_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_tokSfB_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 174080, "bmm_E4m3_E4m3E4m3_Fp32_t128x32x512_s2_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_tokSfB_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 256, "15cdfa203783589d46b84b273c10c7e0848a5fa3cf158ff0d5c39c9ac0c26bd0", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 128 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 2 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 1 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x32x512u2_s2_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_tokSfB_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x32x512u2_s2_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_tokSfB_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 174080, "bmm_E4m3_E4m3E4m3_Fp32_t128x32x512u2_s2_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_tokSfB_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 256, "870c199a8a501a4b629f5d1bf97d1c8d76d06c295d76f33cf010e35cc253e99d", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 128 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 2 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 1 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 115712, "bmm_E4m3_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a", 448, "fe3731381664f21d8643a66720b297fdf4dd2b2e22475bb19149c23c4cb49856", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 115712, "bmm_E4m3_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a", 416, "0a5c5f1765b677f4b12a5a77bd42255d2b44f00d1f79a33656a248d7f1673b5a", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 115712, "bmm_E4m3_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a", 448, "ae778ccd730dc872964386dd1f58bf3aa7c396bfa63f5879f6f1e6723bbfb745", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 115712, "bmm_E4m3_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a", 416, "92968153ee4f2bbc6fee0614b1424ebd9639a9f476ba287ccea6d66931429eab", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 115712, "bmm_E4m3_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a", 448, "b2efeb8d3eb03af3ddbfc5b6a02fcaebc2d6f7d966a53f11ce5b6246c05f1eac", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 115712, "bmm_E4m3_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a", 416, "9f9f0047824e413b32a53d05b28fc137716be2d3fa77718a31fc8762932530b8", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 115712, "bmm_E4m3_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a", 448, "3c0b2b33d7a82d5eeb15bbfca04aa57e0a3d1e124623c87ab7be6923619e0f2d", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 115712, "bmm_E4m3_E4m3E4m3_Fp32_t128x64x128_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a", 416, "c813ada1612f3d3529058eca1a77529917a0cdd30adb9e4f6a023134be4342e1", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 115712, "bmm_E4m3_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a", 448, "0f3a5c9b19f422182c88f22ef643708e3a7608cbc075eb0a64fe286805056570", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 115712, "bmm_E4m3_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a", 416, "c684c20e117ed0c344362db6def8db796da95338da784c470a0a469aceffc1b2", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 115712, "bmm_E4m3_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a", 448, "9ea1b44d1ae9b5625b2c8c7bee3ea564a299441610c6e8c5305691f7972487f8", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 115712, "bmm_E4m3_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a", 416, "1d3d411a31710c47d5820a6ac8b14efd137694b3dcf085e9cc11f79b72dcec3b", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 115712, "bmm_E4m3_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a", 448, "c1352b94785c2f2aaf5b073b4738128a81b7a119d25da73f72745aaa0113bcbe", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 115712, "bmm_E4m3_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a", 416, "9c3ecac762384d250afc9ea993dded8ed1122a529b7ba6f7699683ad9b322e0c", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 115712, "bmm_E4m3_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a", 448, "085bc5bf30ea47141214a0d9c5b166fd6df569cfc14c38ad7d9e44174fc1af0f", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 115712, "bmm_E4m3_E4m3E4m3_Fp32_t128x64x128u2_s4_et64x64_m64x64x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a", 416, "4870220f05c78df32faddc8de55a4edf07991b400bb1b206300106ef72a8cb55", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x64x512_s2_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_tokSfB_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x64x512_s2_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_tokSfB_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 211968, "bmm_E4m3_E4m3E4m3_Fp32_t128x64x512_s2_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_tokSfB_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 256, "8fc4f3a93233412c4dad29628b229a2a32ef9149b992cad845c3bde762d427e3", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 128 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 2 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 1 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x64x512u2_s2_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_tokSfB_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x64x512u2_s2_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_tokSfB_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 211968, "bmm_E4m3_E4m3E4m3_Fp32_t128x64x512u2_s2_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_tokSfB_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 256, "530590ee26d518825cbf6584939b657dcbc1ea07f6c2d89c736f7630ec7aa9bd", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 128 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 2 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 1 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_sm100a_cubin_len, 61440, "bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_sm100a", 384, "d42f2fd34e8cf59c328baa0435afc72534d22185f8808487075df41e56fc09ca", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 1 +, /* mNumTokens */ 0 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a_cubin_len, 61440, "bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a", 384, "ec7648b300c2908dfa94d1ee2140649ede8eaf25c283bbbb0374718c18c50c79", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 1 +, /* mNumTokens */ 0 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_sm100a_cubin_len, 61440, "bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_sm100a", 384, "dc6001cf0e3e3c41a23eb6ae9b929d0bed1f79528225b0dcb3bec818334c9049", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 1 +, /* mNumTokens */ 0 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a_cubin_len, 61440, "bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a", 384, "b624dad0cbb1c4b9fe803dfee5bbd9d2d05c3269d47dd7014031dbddbf836699", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 1 +, /* mNumTokens */ 0 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 149504, "bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a", 448, "b7b0cc6d9c6ae0ebe46a369e14273a1ce80dcbf3a52f9b25e54a7d7a9820598a", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 8 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 149504, "bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a", 416, "aec05f561a21befa59a224f356b959b406fd2c2ea01c6c0cef9e9f884f7edefa", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 8 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 149504, "bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a", 448, "ff0f609512a66261d9cacdd4f6d7c6f0a0df1f9ec6e8b87d45fe7807962cbcf8", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 8 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 149504, "bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a", 416, "47ec62d07f57e6931f53a84a34fc80ba662dd932a740cde606196a18832e5234", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 8 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 149504, "bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a", 448, "e1c21a7380cbe830099603605416a0ce4a3ae04287ce0db38d45202610486425", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 8 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 149504, "bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a", 416, "e9d722f5c6c64c294fc44cb3265d930a5aa1e2db57b6785ca02a2c20b7a53a52", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 8 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 149504, "bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a", 448, "97c72e81f143cb0f5316f3890b2f2efdb4e8e6f96f673392ecc95bc964e3cfb7", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 8 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 149504, "bmm_E4m3_E4m3E4m3_Fp32_t128x8x128_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a", 416, "013a501582343cf23c12e814739f06d0fcc6c46723c016cf862c0dc8446c58f8", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 8 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_sm100a_cubin_len, 61440, "bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_sm100a", 384, "f69a3334f3dd7dbba1b8055012f9f6a7c4716360c4081daccad177d6bdbe1de5", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 1 +, /* mNumTokens */ 0 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a_cubin_len, 61440, "bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a", 384, "26f9cdd0e653cb1d23bfcc62971c095cd9ec1e7b903680e0bcb52295bcf7bf4d", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 1 +, /* mNumTokens */ 0 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_sm100a_cubin_len, 61440, "bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_sm100a", 384, "e7a181f3677b6ab072fe2af99e1a5b7929f78a9562c9d39e9de3c1fe714211a0", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 1 +, /* mNumTokens */ 0 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a_cubin_len, 61440, "bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a", 384, "dddcf3923d544d600a0722153bce5eae2c6ab823a97cbdf291277bf5e01c4ab0", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 1 +, /* mNumTokens */ 0 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 149504, "bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a", 448, "46c66d9f32de8f08c6db20bcdd8744ba8fdeadc37898b2c5352fa6db9ebbd2db", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 8 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 149504, "bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a", 416, "eac88b97c198140559ef278cafa4cf63575292e7b35346416b6e78cb1c9bb5c5", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 8 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 149504, "bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a", 448, "9c96ccc014badff11a1f50c3c40f5b299ad03fd5ba8a32c8e3bc42f490dba266", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 8 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 149504, "bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a", 416, "60d81ae7da341580c6ba652f6c5b9c25c9c0cd3c55bb510e6641f923b7b3633e", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 8 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 149504, "bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a", 448, "9daf7ce62048baf89dcc9b28d09fb98b4033ce57f25654a95e79ec819e1d84fb", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 8 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 149504, "bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a", 416, "bfbe62be4152bd934ee6553669dd7f59e211f139770ab4b64c21882e3f52c270", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 8 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 149504, "bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedP_bN_ldgsts_clmp_dynBatch_sm100a", 448, "52fe80baf1b55bfac15ca876b5c2b4111d1f475fa176219ba491543cbe908360", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 8 +, /* mNumStagesMma */ 4 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a_cubin_len, 149504, "bmm_E4m3_E4m3E4m3_Fp32_t128x8x128u2_s8_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_ldgsts_clmp_dynBatch_sm100a", 416, "bc2ffee4831f12acf4e54ac3ef3f6279c5517100ee04d16167587114ecabf409", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 8 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a_cubin_len, 216064, "bmm_E4m3_E4m3E4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a", 224, "6329f7ec972e432372452c9988851e27dffa49e10b922723e5bef3f9540b2987", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 1 +, /* mNumTokens */ 0 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_tokSfB_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_tokSfB_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 216064, "bmm_E4m3_E4m3E4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_tokSfB_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 256, "54364b012c0467825b7a02615275ac53a6131027fbe10d78589c2211af6bd4b0", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 1 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a_cubin_len, 216064, "bmm_E4m3_E4m3E4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a", 224, "13992f122109be6fdfc576480470f0a0a8a56c6bfa2210b317debac5045e1fcf", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 1 +, /* mNumTokens */ 0 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_E4m3E4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_tokSfB_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_E4m3_E4m3E4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_tokSfB_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 216064, "bmm_E4m3_E4m3E4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_tokSfB_schedS_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 256, "c3b553e55fb4e898d86ee14ff1bb37ec1a6a3468d5f337e0fc369667a7da4218", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 1 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x16x256_s5_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x16x256_s5_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 199680, "bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x16x256_s5_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 352, "f1a7ea01b8739ed28f5a7e5a3c68715b0feeed07b8a372ee76b024c32e454481", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 5 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x16x256_s5_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x16x256_s5_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 199680, "bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x16x256_s5_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 320, "8f87812e830b0ef68528fb1185fcd396369825795664fedbb01903d14aa1b45a", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 5 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x16x256u2_s5_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x16x256u2_s5_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 199680, "bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x16x256u2_s5_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 352, "a2e4f42148f636846ccf23df49547a8624cd6390973c265165222271a15a7839", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 5 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x16x256u2_s5_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x16x256u2_s5_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 199680, "bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x16x256u2_s5_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 320, "66b38fcde671df041618bb144ad0ac9dca3af15338960c3d72f32763b0f46f9b", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 5 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x32x256_s5_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x32x256_s5_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 222208, "bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x32x256_s5_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 352, "14c49c0ba19d98a5daeac0f875923a81f1280a2aa997ee2b17532588118fd747", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 5 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x32x256_s5_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x32x256_s5_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 222208, "bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x32x256_s5_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 320, "206d9ef92877403581df57ff9dafbd8846022f42d05460a6999bfbd3c13346f5", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 5 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x32x256u2_s5_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x32x256u2_s5_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 222208, "bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x32x256u2_s5_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 352, "cf4a9445827af25ad5509c977520e71d240c1cf481218c5a4ce8f42e4365497b", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 5 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x32x256u2_s5_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x32x256u2_s5_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 222208, "bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x32x256u2_s5_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 320, "6638a120534f8e5057d418a804d85ed31d792805a765dcf31007d2d544e59a58", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 5 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x64x256_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x64x256_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 217088, "bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x64x256_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 352, "ba3623bd75c6d2886dcb8d909ea4d22aebfcb818322f82d31112b1f0b4c788d6", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x64x256_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x64x256_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 217088, "bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x64x256_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 320, "fbbc609f53caa9b124fb5a57fb7e0a0eaebe8bb331ca20f17e8906627ac31c1d", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x64x256u2_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x64x256u2_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 217088, "bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x64x256u2_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 352, "5ff1d8d57151079800893c8a7349faa3aa489e6c1a01c197ded53a390d3af1fc", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x64x256u2_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x64x256u2_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 217088, "bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x64x256u2_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 320, "c04da5f63259b901c8aaed6749707268f14f784e923908e5555a7d6042f8793e", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 188416, "bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 352, "ebc4029975da6fc6d2ac7c907a89c1a1e06aae8136ee6ccd9ae1ed86e7fc26a2", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 5 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 188416, "bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 320, "ea5395a0463c0a39e6dae7be45d4080c0695fc3edc6678f813a6769c2508eb16", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 5 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 188416, "bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 352, "9bdac69520fd1a4c1915b0612ceb733cc58653e21c749af4cfdb42a5fe08e8e4", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 5 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 188416, "bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 320, "71c50c6ccb8be1c7cde7c18375e561f4463133a041691f8592e1fecc67d2d7e4", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 5 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 224256, "bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 352, "6c8b4c915bb3b8131301c4ca922e96c5d5af6efcca0beba70a43b4ea08a887db", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 224256, "bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 320, "5581245c8c20d8357f33d296c4d3448143d62406575bc5ccfa175e133ec32dfa", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 224256, "bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 352, "9609ea9566913c76e84a837b7511ec63b835af3653a3e41ca4c0f591fc661b44", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 224256, "bmm_E4m3_MxE2m1E4m3_castMxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 320, "2a2184a71fcc77e60a4a7686a152bf8eb660c5da27df105eb786915ea670d2e0", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Fp16_E2m1E2m1_Fp32_t128x8x512_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a_cubin, Bmm_Fp16_E2m1E2m1_Fp32_t128x8x512_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a_cubin_len, 166912, "bmm_Fp16_E2m1E2m1_Fp32_t128x8x512_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a", 448, "6d28d0eec821539c8790a3f998ff68987199a3c34d8d3775ac653289a8bab2da", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826818) +, /* mDtypeB */ trtllm::gen::Dtype(17826818) +, /* mDtypeC */ trtllm::gen::Dtype(1052679) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17826818) +, /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaKind */ trtllm::gen::MmaKind(4) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 1 +, /* mNumTokens */ 0 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Fp16_E2m1E2m1_Fp32_t128x8x512u2_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a_cubin, Bmm_Fp16_E2m1E2m1_Fp32_t128x8x512u2_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a_cubin_len, 166912, "bmm_Fp16_E2m1E2m1_Fp32_t128x8x512u2_s4_et128x8_m128x8x64_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a", 448, "ee70583d209fdd0267dbcadfed212c7a11cf263c63ed566fe27d479d1b7f5d21", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826818) +, /* mDtypeB */ trtllm::gen::Dtype(17826818) +, /* mDtypeC */ trtllm::gen::Dtype(1052679) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826818) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17826818) +, /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 64 +, /* mMmaKind */ trtllm::gen::MmaKind(4) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(1) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 1 +, /* mNumTokens */ 0 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Fp16_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_sm100a_cubin, Bmm_Fp16_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_sm100a_cubin_len, 61440, "bmm_Fp16_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_sm100a", 384, "0c5f5ffa8eb28913c87e78e5bf259dc4dce206544479aec4235750e74ac11835", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052679) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 1 +, /* mNumTokens */ 0 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Fp16_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a_cubin, Bmm_Fp16_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a_cubin_len, 61440, "bmm_Fp16_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a", 384, "98c29f7ebdc1e95425a63da5002644a84dd4ff5b0c9d9c3222b405a3768c6806", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052679) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 1 +, /* mNumTokens */ 0 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Fp16_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_sm100a_cubin, Bmm_Fp16_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_sm100a_cubin_len, 61440, "bmm_Fp16_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_sm100a", 384, "d140f74147be90c9ccb65d49f231dde181e34e3742158a7a25700617974bfa89", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052679) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 1 +, /* mNumTokens */ 0 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Fp16_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a_cubin, Bmm_Fp16_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a_cubin_len, 61440, "bmm_Fp16_E4m3E4m3_Fp32_t128x8x128_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a", 384, "ddf6ebc5fd61ee1618e7a846e2010d90b8f97e940b9d15fb1a4013c2f77ee92f", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052679) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 1 +, /* mNumTokens */ 0 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Fp16_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_sm100a_cubin, Bmm_Fp16_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_sm100a_cubin_len, 61440, "bmm_Fp16_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_dsFp8_schedS_bN_clmp_sm100a", 384, "b942baa865cf92a34a6cb05de9f87fdd099a64d5240634619c164207bb0f6ce4", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052679) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 1 +, /* mNumTokens */ 0 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Fp16_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a_cubin, Bmm_Fp16_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a_cubin_len, 61440, "bmm_Fp16_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_BN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a", 384, "2c03409fe9e5a05e92ba863cb541d81c150593d8da22c72987e1d8ab97b7f037", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ 128 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052679) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(2) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 1 +, /* mNumTokens */ 0 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Fp16_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_sm100a_cubin, Bmm_Fp16_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_sm100a_cubin_len, 61440, "bmm_Fp16_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_dsFp8_schedS_bN_clmp_sm100a", 384, "df6216b63156119ec1904594120b1b27a8e80b24930691c4cfaa55fe85fcffde", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052679) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 1 +, /* mNumTokens */ 0 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Fp16_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a_cubin, Bmm_Fp16_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a_cubin_len, 61440, "bmm_Fp16_E4m3E4m3_Fp32_t128x8x128u2_s3_et64x8_m64x8x32_cga1x1x1_16dp256b_TN_transOut_noShflA_dsFp8_schedS_bN_clmp_sm100a", 384, "7919353c891ab408893abc7298e876421fd8c59eed469ed98e31124752b1c931", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052679) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 64 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 1 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 64 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 2 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 0 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 128 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 1 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 1 +, /* mNumTokens */ 0 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Fp16_E4m3E4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a_cubin, Bmm_Fp16_E4m3E4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a_cubin_len, 217088, "bmm_Fp16_E4m3E4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a", 224, "42552b1145a7637fe53aeb1b91fad1cfb07668f59439e9996507f3079eb475c8", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052679) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 1 +, /* mNumTokens */ 0 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_Fp16_E4m3E4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a_cubin, Bmm_Fp16_E4m3E4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a_cubin_len, 217088, "bmm_Fp16_E4m3E4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_bN_clmp_sm100a", 224, "194aa97105e1648571be4a622ec13fe9b5fed84902680eb7450be52b3a1c5fe6", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(0) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(1050629) +, /* mDtypeB */ trtllm::gen::Dtype(1050629) +, /* mDtypeC */ trtllm::gen::Dtype(1052679) +, /* mDtypeMmaA */ trtllm::gen::Dtype(1050629) +, /* mDtypeMmaB */ trtllm::gen::Dtype(1050629) +, /* mEnablesEarlyExit */ 0 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(2) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(3) +, /* mSfLayoutC */ trtllm::gen::SfLayout(3) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 1 +, /* mNumTokens */ 0 +, /* mRouteImpl */ batchedGemm::RouteImpl(0) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 0 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x16x256_s3_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x16x256_s3_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 123904, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x16x256_s3_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 512, "78505ef3c73f04a0d53445d639ff707704ee4c6d886831736c14ecec82600efe", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x16x256_s3_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x16x256_s3_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 123904, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x16x256_s3_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 480, "599659954bd80cfda041f664d83c13ae89c605f0cbda94cada70f00a310bb2ea", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x16x256_s4_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x16x256_s4_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 161792, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x16x256_s4_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 512, "6db5f9d10464bade0417c0fc45dccc499cd8cb1283aae6a36d2b902531a484c3", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x16x256_s4_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x16x256_s4_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 161792, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x16x256_s4_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 480, "daa5e40fe2a0fd55c17306459dfd56952c665863ffdb544b2e7160040ca210f5", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x16x256u2_s3_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x16x256u2_s3_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 123904, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x16x256u2_s3_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 512, "de84c185f683a95f68b3f0628f43d6ace368189f7534677f34df4bdb23ca65fa", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x16x256u2_s3_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x16x256u2_s3_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 123904, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x16x256u2_s3_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 480, "a16e7157ebc69edc07507cb56255930a104b86096e923e31252c28179540f337", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x16x256u2_s4_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x16x256u2_s4_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 161792, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x16x256u2_s4_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 512, "186dce0ee0341e4f90ce0e673478e7041fbb57f0338d8f9d90dd255d4622820f", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x16x256u2_s4_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x16x256u2_s4_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 161792, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x16x256u2_s4_et128x16_m128x16x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 480, "91f75f1bbb0773becd4cc8f2fe187ba67af9aa9c640d6adc7036803c4c882fe7", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 16 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 16 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 16 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x32x256_s3_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x32x256_s3_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 138240, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x32x256_s3_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 512, "4bf94cf88b3e79a1232408d75aa723d1c8a4d5e6cee60bbd6dce06fdf8b67bdc", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x32x256_s3_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x32x256_s3_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 138240, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x32x256_s3_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 480, "a7646e3bd92185c13588d0238802f633e04a171c95c14a134b1786b3c9bca4fb", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x32x256_s4_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x32x256_s4_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 181248, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x32x256_s4_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 512, "76fd23564cb7223eacd20c08df3aace2866882fc4a0edda811c6db79532b653b", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x32x256_s4_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x32x256_s4_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 181248, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x32x256_s4_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 480, "613d320526ab9f209870c363945d2815213cc0238e1ba5c448e3459b7e7bd531", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x32x256u2_s3_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x32x256u2_s3_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 138240, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x32x256u2_s3_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 512, "be616bf5cd5656a26d2304d2359c1cc82a5eb750a30298b8420f50089f2bcf23", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x32x256u2_s3_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x32x256u2_s3_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 138240, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x32x256u2_s3_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 480, "46d660330199d6395de999cfa7734fc16791d535777b54d565e68424d62ca10b", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x32x256u2_s4_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x32x256u2_s4_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 181248, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x32x256u2_s4_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 512, "f8a6d222dc9affb478e6dbb2ac57b61e607af8200225f9d7e433f74cfe2fd289", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x32x256u2_s4_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x32x256u2_s4_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 181248, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x32x256u2_s4_et128x32_m128x32x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 480, "fcbc96fd294b7fdd364c235b03cb2e04c151573a0564f3dbb28a407c707c4377", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 32 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 32 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 32 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x64x256_s3_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x64x256_s3_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 168960, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x64x256_s3_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 512, "786f358bf1789989723720d0d932ec12853e49d019db3b09be5426cdeca358f0", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x64x256_s3_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x64x256_s3_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 168960, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x64x256_s3_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 480, "5804ce859c989fa2266dfa1d550f5a6c0df19b28d7ca6f9814980424cdb81d09", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x64x256_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x64x256_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 220160, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x64x256_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 512, "ab64018790ccb1f457648cb75296ec7d4d091d54350cd17bc7e5e468a0cb3316", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x64x256_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x64x256_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 220160, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x64x256_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 480, "e7f67ed99ebf367b53811ecb55579173cf697806f420ed85a80151d056c23415", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x64x256u2_s3_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x64x256u2_s3_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 168960, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x64x256u2_s3_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 512, "c7ed8f6d742ea142fda745c5715ba9dcbf8941a343af271d7844d55f02feeb57", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x64x256u2_s3_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x64x256u2_s3_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 168960, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x64x256u2_s3_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 480, "2cee9f40f0ed9d515faabb3a8d0352b74103bfb56bcbff14674f620527fe50fb", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x64x256u2_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x64x256u2_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 220160, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x64x256u2_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 512, "17c89d1bb74c646281024f4348eaff386b34dcda2ec28b4df28d52c53e825cfe", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x64x256u2_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x64x256u2_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 220160, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x64x256u2_s4_et128x64_m128x64x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 480, "fc0af5af884f070d4f82fd35292f2433fd4904594af766f3db1bba3961f9cf5f", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 64 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 64 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 64 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s4_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s4_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 152576, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s4_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 512, "ffba61927b3ea1e1b7b29084f89e33dd5a6e83fe32d2b3e79eddb3446da2a761", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s4_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s4_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 152576, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s4_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 480, "49824cbed86455cc7836b05c0b5261ee51e94ecee498b37f64d1cc923e9a43c9", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s4_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s4_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 159744, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s4_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 512, "fda32440ec4d9a6c98afd95404e63e22274f1b3ab4ff7744cfb6ba354379aa93", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 2 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 2 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(2) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s4_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s4_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 151552, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s4_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 480, "a60d72364a4180f4599e443f62fea83f37830df9e328a34d3fae7e1f825a0a94", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 2 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 2 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(2) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 188416, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 512, "3bcb4529126bf44527c333e75a1b45122e2e141abba802146be75c524c4a8e57", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 5 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 188416, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 480, "afd00d9e20e5f50fa2d83bc2d4179582b100c6faf73fb627e1334f049fde0ce0", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 5 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 195584, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 512, "10b2e0624341b0f3b28d01764245248f5a942553d82018cdbf2d5d0655f09aad", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 2 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 2 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 5 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(2) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 187392, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s5_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 480, "87b861e015c9409b2aff171d508a02308772948e4c36afdc9c9053dd699bd940", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 2 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 2 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 5 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(2) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s6_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s6_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 224256, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s6_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 512, "26abed2f58502b931d728d01683291c0f7a5415b275e7b476db2397654ba5a02", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s6_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s6_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 224256, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s6_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 480, "ead5254e30dfe83e246606071c686200dec7f40e4f3e72119b4cf42d975e5e40", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s6_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s6_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 231424, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s6_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 512, "c2467d2a97ac86edb03b061d148702df24715febb02edb7aa371c694eef8a37a", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 2 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 2 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(2) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s6_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s6_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 223232, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256_s6_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 480, "3a197c4d4b40880ce4a6b9c309c38fa36f2df1c5c02a0b1199aba2d560fc8a50", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 2 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 2 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(2) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s4_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s4_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 152576, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s4_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 512, "7584dabe78ad0dc6b1b627e7eec0d21c91c4b9c1cd53eff3240b990d10e552ca", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s4_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s4_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 152576, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s4_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 480, "71dda7f91b6925618277b436bdb95c1c609c89d25f6c9227b0c421dfdd8f0382", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s4_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s4_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 159744, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s4_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 512, "2542a115ddfab3bbcb8816a2b4af26cc2d881f2c6d6552bd7b4200b40ff4b60c", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 2 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 2 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(2) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s4_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s4_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 151552, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s4_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 480, "6283f013fe3b5444f423f42a61ad151f7beae9145748ced3caa0674b0675c3d4", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 2 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 2 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 4 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(2) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 188416, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 512, "acc297829bcfd0e45ced4300cf69edeed6e727a32eda1d067da6403747784e3e", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 5 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 188416, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 480, "92139dfe591eaca228b2f79721e981679803376aa692a3366e351f94ac51ede7", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 5 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 195584, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 512, "5a86ddfe2f7b573941b49fe1e99497fa10b302fd96d3dc4dde0304311912ad7e", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 2 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 2 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 5 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(2) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 187392, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s5_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 480, "d869e97fc378587d8f55934e4118be427c5c31a2156dbafdce2e9a4931e20848", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 2 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 2 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 5 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(2) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s6_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s6_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 224256, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s6_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 512, "f8cba023c6150b5e80771d19703d377e1ef61ba0ddaa80bfbcd7a4b52d762bbe", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s6_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s6_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 224256, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s6_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 480, "1de7b15c33a7b5dfb87cde99a2693f6bf2ece88f3706c94b3507413ada41aa60", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s6_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s6_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 231424, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s6_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 512, "8d3c3428f4acb8c1d45109f2d26df94f8112142689c7ea37e369628ae6215463", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 2 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 2 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(2) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s6_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s6_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 223232, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x256u2_s6_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 480, "6d2b7221be879c4e19b857ef5099521a111254cfdc1211d6415f0bc2c0ae152f", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 2 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 2 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 6 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(2) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 256 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 224256, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 512, "66059be69e378daef40ad4f19a17ac1d9dd5305a589f963a408a88080ad8bcb6", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 224256, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 480, "77f8dab69fc0c346e977b67408f16749de76ed4df80c438f8af37d2ad610befa", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 231424, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 512, "267bfa7448c730809d49cc3707168e1b1f9d4a7824b8449af51ef0e1bcd6ef68", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 2 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 2 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(2) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 223232, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x512_s3_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 480, "91b2af33c93bf7e89afca5ea4b8c58b99db79ed713f550d2e31243cce072d5e0", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 2 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 2 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(2) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 0 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 224256, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 512, "713657671934397700abf3f77cdd20da3df2cf9cc9eaebff78bc2cd5763bf866", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 224256, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x1_16dp256b_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 480, "6cc8de91f83ec44751953e155beee55a2bcbfaf6c0fc4a1256e527466484c61a", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 1 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 1 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(0) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 231424, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedP_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 512, "0321ef95643ff3da86a79a505034de18b1fc0d9f7f6e287b17e82d925322569a", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 2 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 2 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 2 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 2 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(2) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(1) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +{Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin, Bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a_cubin_len, 223232, "bmm_MxE4m3_MxE2m1MxE4m3_Fp32_t128x8x512u2_s3_et128x8_m128x8x32_cga1x1x2_16dp256b_splitK2_TN_transOut_schedS_biasM_bN_ldgsts_clmp_swiGlu_dynBatch_sm100a", 480, "4380454862dfce06463093bc1feead0c827a31e2080b33cf6bf9988210e6c6eb", { /* mAllReduceAlgo */ gemm::AllReduceAlgo(0) +, /* mBiasType */ gemm::BiasType(1) +, /* mBlockK */ -1 +, /* mClusterDimX */ 1 +, /* mClusterDimY */ 1 +, /* mClusterDimZ */ 2 +, /* mDtypeAcc */ trtllm::gen::Dtype(1056776) +, /* mDtypeA */ trtllm::gen::Dtype(17826828) +, /* mDtypeB */ trtllm::gen::Dtype(17827853) +, /* mDtypeC */ trtllm::gen::Dtype(17827853) +, /* mDtypeMmaA */ trtllm::gen::Dtype(17826828) +, /* mDtypeMmaB */ trtllm::gen::Dtype(17827853) +, /* mEnablesEarlyExit */ 1 +, /* mEnablesDelayedEarlyExit */ 0 +, /* mEnablesGlobalPtxKnobs */ 1 +, /* mEpilogueLdtmDps */ 16 +, /* mEpilogueLdtmBits */ 256 +, /* mEpilogueTileM */ 128 +, /* mEpilogueTileN */ 8 +, /* mGridTriggerSecondaryA */ 0 +, /* mGridTriggerSecondaryB */ 1 +, /* mGridWaitForPrimaryEarlyExit */ 1 +, /* mGridWaitForPrimaryA */ 0 +, /* mGridWaitForPrimaryB */ 1 +, /* mHoistLoadTaskInit */ 1 +, /* mHoistMmaTaskTryWaits */ 0 +, /* mK */ 2048 +, /* mKernelTraits */ {} +, /* mLayoutA */ gemm::MatrixLayout(0) +, /* mLayoutB */ gemm::MatrixLayout(0) +, /* mM */ 256 +, /* mMmaK */ 32 +, /* mMmaKind */ trtllm::gen::MmaKind(5) +, /* mMmaM */ 128 +, /* mMmaN */ 8 +, /* mMockAllReduce */ 0 +, /* mN */ 256 +, /* mNumSlicesForSplitK */ 2 +, /* mNumSlicesForSliceK */ 1 +, /* mNumStages */ 3 +, /* mNumStagesMma */ 1 +, /* mNumStagesMmaWithinWorkTile */ 1 +, /* mNumStagesMmaAcrossWorkTile */ 1 +, /* mNumStagesWorkId */ 3 +, /* mOutputDebugTensors */ 0 +, /* mPatchF2fp */ 0 +, /* mUseShuffledMatrixA */ 1 +, /* mSliceK */ 0 +, /* mSplitK */ gemm::SplitK(2) +, /* mTransposeMmaOutput */ 1 +, /* mTileM */ 128 +, /* mTileN */ 8 +, /* mTileK */ 512 +, /* mUseUnrollLoop2xForMma */ 1 +, /* mUseCustomMmaSchedule */ 1 +, /* mUseHoistTryWaitForCustomMmaSchedule */ 0 +, /* mUseDeepSeekFp8 */ 0 +, /* mUsePerTokenSfA */ 0 +, /* mUsePerTokenSfB */ 0 +, /* mUseTmaStore */ 1 +, /* mUseTwoTmaLoadWarps */ 1 +, /* mUseTwoMmaWarps */ 0 +, /* mSfLayoutA */ trtllm::gen::SfLayout(3) +, /* mSfLayoutB */ trtllm::gen::SfLayout(0) +, /* mSfLayoutC */ trtllm::gen::SfLayout(1) +, /* mSfReshapeFactor */ 1 +, /* mTileScheduler */ gemm::TileScheduler(0) +, /* mActType */ gemmGatedAct::ActType(0) +, /* mClampBeforeAct */ 1 +, /* mBatchedM */ {} +, /* mBatchedN */ {} +, /* mBatchMode */ batchedGemm::BatchedGemmOptions::BatchMode(1) +, /* mNumBatches */ 2 +, /* mIsStaticBatch */ 0 +, /* mNumTokens */ 2 +, /* mRouteImpl */ batchedGemm::RouteImpl(1) +, /* mGridWaitForPrimaryRouting */ 1 +, /* mFusedAct */ 1 +, /* mNumRegsPerThreadNonEpilogueWarp */ 0 +, /* mNumRegsPerThreadEpilogueWarp */ 0 +, /* mNumRegsCastAWarps */ 0 +, /* mUseTmaOobOpt */ 0 + }, gemm::SmVersion::Sm100a}, +#endif // EXCLUDE_SM_100 +}; +// clang-format on +} // namespace kernels +} // namespace tensorrt_llm +} // namespace batchedGemm diff --git a/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/KernelParams.h b/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/KernelParams.h new file mode 100644 index 0000000000000000000000000000000000000000..91d21ea239fa411e7a8f0084f8557e0d365c8e77 --- /dev/null +++ b/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/KernelParams.h @@ -0,0 +1,579 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & + * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include + +#include "BatchedGemmEnums.h" +#include "Enums.h" +#include "TmaDescriptor.h" +#include "trtllm/gen/CommonUtils.h" +#include "trtllm/gen/SfLayoutDecl.h" + +// NOTE: keep this code dependency free. It has to be included by the device code and has to be +// compilable with NVRTC. +#include "KernelParamsDecl.h" + +namespace batchedGemm { + +namespace batchedGemm { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// TODO: Find a better header to put this in, that we can include from here. +template +inline T ceilDiv(T m, T n) { + return (m + n - T(1)) / n; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace tg = trtllm::gen; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace KernelParamsSetup { +#ifdef TLLM_ENABLE_CUDA + +enum class MatrixType { MatrixA = 0, MatrixB, MatrixC }; + +////////////////////////////////////////////////////////////////////////////////////////////////// +// +// Utility functions. +// +////////////////////////////////////////////////////////////////////////////////////////////////// + +template +bool useTmaOobOptA(BatchedGemmOptions const& options) { + return options.mBatchMode == BatchedGemmOptions::BatchMode::BatchM && + doesRouteImplUseNoRoute(options.mRouteImpl) && options.mUseTmaOobOpt; +} + +////////////////////////////////////////////////////////////////////////////////////////////////// + +template +bool useTmaOobOptB(BatchedGemmOptions const& options) { + return options.mBatchMode == BatchedGemmOptions::BatchMode::BatchN && + doesRouteImplUseNoRoute(options.mRouteImpl) && options.mUseTmaOobOpt; +} + +////////////////////////////////////////////////////////////////////////////////////////////////// + +template +bool useTmaOobOptC(BatchedGemmOptions const& options) { + return options.mUseTmaStore && options.mUseTmaOobOpt; +} + +////////////////////////////////////////////////////////////////////////////////////////////////// + +// Create the TMA shape/stride for A/B/C. +template +static auto makeTmaShapeStrideAbc(GemmOptions const& options, int mM, int mN, int mK, int tileM, + int tileN, int tileK, MatrixType matrixType) { + // Weights matrix is A if we transpose the output of MMA (to have it M-major). + // Otherwise, it is B, when the output of MMA is K-major. + bool const isWeights = (matrixType == MatrixType::MatrixA && options.mTransposeMmaOutput) || + (matrixType == MatrixType::MatrixB && !options.mTransposeMmaOutput); + + // Whether to use TMA OOB trick to block out padded dummy tokens and saving BW whenever no routing + // is involved. It applies to batchM and matrixA, or batchN and matrixB, or any case for matrixC. + bool const useTmaOobOpt = matrixType == MatrixType::MatrixA ? useTmaOobOptA(options) + : matrixType == MatrixType::MatrixB ? useTmaOobOptB(options) + : matrixType == MatrixType::MatrixC ? useTmaOobOptC(options) + : false; + + // The outer dimension. + auto numTokens = + (matrixType == MatrixType::MatrixA || matrixType == MatrixType::MatrixC) ? mM : mN; + // The outer dimension tile size. + auto ctaTileNumTokens = + (matrixType == MatrixType::MatrixA || matrixType == MatrixType::MatrixC) ? tileM : tileN; + // The outer dimension of TMA box shape. + auto tileNumTokens = + (matrixType == MatrixType::MatrixC) ? options.mEpilogueTileM : ctaTileNumTokens; + + // The inner dimension. + auto hiddenSize = (matrixType == MatrixType::MatrixC) ? mN : mK; + // The inner dimension tile size. + auto ctaTileHiddenSize = (matrixType == MatrixType::MatrixC) ? tileN : tileK; + // The inner dimension of TMA box shape. + auto tileHiddenSize = + (matrixType == MatrixType::MatrixC) ? options.mEpilogueTileN : ctaTileHiddenSize; + + // Swap matrix C sizes if output is transposed. + if (matrixType == MatrixType::MatrixC && options.mTransposeMmaOutput) { + std::swap(numTokens, hiddenSize); + std::swap(ctaTileNumTokens, ctaTileHiddenSize); + std::swap(tileNumTokens, tileHiddenSize); + } + + // For a fused activation kernel, the hidden size of output is halved. TODO: That's true for + // gated activations but not regular activations. + if (options.mFusedAct && matrixType == MatrixType::MatrixC) { + hiddenSize /= 2; + tileHiddenSize /= 2; + ctaTileHiddenSize /= 2; + } + + // The cute tensor shape for A/B: (numTokens, hiddenSize). + // Note that TMA descriptor expects the first dimension's stride to be + // 1, so swap the first two dimension so that the hiddenSize dimension comes first. + + // Activations matrix is 2D (sum(divUpMul(M[bi], tileM) for bi in B), K). + std::vector shape = {static_cast(hiddenSize), + static_cast(numTokens)}; + if (useTmaOobOpt /* also implies input/output activation */) { + // If TMA OOB optimization is used, we use 3D logical shape (M, tileM, K) or (N, tileN, K). + // The outer dimension is extended to make room for the possible counterbalance positive + // offset from the middle "bound" dimension. The counterbalance should be no more than + // ctaTileNumTokens. + shape = {static_cast(hiddenSize), static_cast(ctaTileNumTokens), + static_cast(numTokens + ctaTileNumTokens)}; + } else if (isWeights) { + // If the matrix is a weights matrix, we use 3D logical shape (B, M, K) or (B, N, K). + shape = {static_cast(hiddenSize), static_cast(numTokens), + static_cast(options.mNumBatches)}; + } + + // Assemble the stride (strideTokens, 1). + // Swap the first two dimension as mentioned before. + std::vector stride = {1, static_cast(hiddenSize)}; + if (useTmaOobOpt) { + stride = {1, static_cast(hiddenSize), static_cast(hiddenSize)}; + } else if (isWeights) { + stride = {1, static_cast(hiddenSize), + static_cast(hiddenSize) * static_cast(numTokens)}; + } + + // Assemble the box shape + std::vector tileShape = {tileHiddenSize, tileNumTokens}; + + // Alternate layouts (MajorMn and BlockMajorK) do not apply to matrixC + if (matrixType != MatrixType::MatrixC) { + gemm::MatrixLayout layout = + (matrixType == MatrixType::MatrixA) ? options.mLayoutA : options.mLayoutB; + // Note, only the weights support non MajorK layouts + if (layout == gemm::MatrixLayout::MajorMn) { + // Apply transpose if necessary + std::swap(shape[0], shape[1]); + stride[1] = numTokens; + std::swap(tileShape[0], tileShape[1]); + } else if (layout == gemm::MatrixLayout::BlockMajorK) { + // Set shapes based on blocking layout + shape = {static_cast(options.mBlockK), static_cast(numTokens), + static_cast(mK / options.mBlockK), + static_cast(options.mNumBatches)}; + stride = {1, static_cast(options.mBlockK), + static_cast(numTokens * options.mBlockK), + static_cast(hiddenSize * numTokens)}; + + // If blockK > tileK, then the inner most box size will be based on the tile + int32_t const tileBlockK = std::min(options.mBlockK, tileHiddenSize); + tileShape = {tileBlockK, tileNumTokens, tileHiddenSize / tileBlockK}; + } + } + + return std::make_tuple(shape, stride, tileShape); +} + +// Create the TMA shape/stride for A/B block scaling factors. +static auto makeTmaShapeStrideSfAb(int mM, int mN, int mK, MatrixType matrixType, int tileM, + int tileN, int tileK, tg::Dtype dtypeElt, tg::SfLayout layout, + int sfReshapeFactor) { + // The outer dimension. + auto numTokens = matrixType == MatrixType::MatrixA ? mM : mN; + // The inner dimension. + auto hiddenSize = mK; + // The outer tile dimension. + auto numTokensPerTile = matrixType == MatrixType::MatrixA ? tileM : tileN; + // The inner tile dimension. + auto hiddenSizePerTile = tileK; + // Number of elements per scaling factor. + const int32_t numEltsPerSf = (dtypeElt == tg::Dtype::E2m1) ? 16 : 32; + + switch (layout) { + case tg::SfLayout::R128c4: { + // The scaling factor tensor packs 128x4 tiles into contiguous 512B blocks. + // The 512B block maps to a 32x16B (32x128b) block in TMEM. + // See https://nvbugspro.nvidia.com/bug/4165523 + // + // Additionally, we have to meet constraints of TMA that the box dimensions are less + // than 256 and boxDim[0] is a multiple of 16B. + // + // The "logical" tensor is: [outer, inner / numEltsPerSf] + // The aforementioned format is: [outer / 128, inner / numEltsPerSf / 4, 512] + // The shape we use for TMA is: [outer / 128, inner / numEltsPerSf / 4, 2, 256] + + auto shape = std::vector{ + 256, 2, static_cast(ceilDiv(hiddenSize, numEltsPerSf * 4)), + static_cast(ceilDiv(numTokens, 128))}; + + std::vector stride(shape.size()); + stride[0] = 1; + for (size_t i = 1; i < shape.size(); i++) { + stride[i] = shape[i - 1] * stride[i - 1]; + } + + auto tileShapes = std::vector{ + 256, 2, static_cast(ceilDiv(hiddenSizePerTile, numEltsPerSf * 4)), + static_cast(ceilDiv(numTokensPerTile, 128))}; + + return std::make_tuple(shape, stride, tileShapes); + } + + case tg::SfLayout::R8c4: { + // The scaling factor tensor packs 8x4 tiles into contiguous 32B blocks. + // + // As the inner dimension (k) is often a multiple of the tile size, we can reshape to use + // fewer read requests, if the tile dimensions allow. It does not reduce the number of + // instructions. + // + // I.e., let's define r = min(⌈hiddenSizePerTile / (numEltsPerSf * 4)⌉, 8) + // + // The "logical" tensor is: [outer, inner / numEltsPerSf] + // The 8x4 SF layout is: [⌈outer / 8⌉, inner / (4 * numEltsPerSf), 32] + // The TMA tensor shape is: [⌈outer / 8⌉, inner / (4 * numEltsPerSf * r), r * 32] + // + // The caveat of NumRepeats>1 is we must pad the hidden dimension of SF to multiples of + // NumRepeats * numEltsPerSf * 4. + + // Detect if the supplied factor is power of 2. E.g., 0b0100 and (0b0100 - 1) == 0b0000. + int const r = sfReshapeFactor; + if (r > 0 && (r & (r - 1)) != 0) { + throw std::runtime_error("mSfReshapeFactor must be positive and a power of 2. Found " + + std::to_string(r)); + } + + // Sanitize number of repeats so it doesn't exceed the dimension. + int const repeats = std::min(ceilDiv(hiddenSizePerTile, numEltsPerSf * 4), r); + + // Detect if the input hidden size K is a multiple of the repeats. + if (ceilDiv(hiddenSize, numEltsPerSf * 4) % repeats != 0) { + throw std::runtime_error( + "SF hiddenSize K (" + std::to_string(ceilDiv(hiddenSize, numEltsPerSf * 4)) + + ") must be a multiple of repeats (" + std::to_string(repeats) + ")"); + } + + auto shape = std::vector{ + static_cast(repeats * 32), + static_cast(ceilDiv(hiddenSize, numEltsPerSf * 4 * repeats)), + static_cast(ceilDiv(numTokens, 8))}; + + std::vector stride(shape.size()); + stride[0] = 1; + for (size_t i = 1; i < shape.size(); i++) { + stride[i] = shape[i - 1] * stride[i - 1]; + } + + auto tileShapes = std::vector{ + static_cast(repeats * 32), + static_cast(ceilDiv(hiddenSizePerTile, numEltsPerSf * 4 * repeats)), + static_cast(ceilDiv(numTokensPerTile, 8))}; + + return std::make_tuple(shape, stride, tileShapes); + } + + default: + throw std::runtime_error("Unsupported SF layout"); + } + return std::make_tuple(std::vector{}, std::vector{}, std::vector{}); +} +template +static KernelParams setKernelParams( + GemmOptions_ const& options, bool const batchM, void const* ptrA, void const* ptrB, void* ptrC, + void const* dSfA, void const* dSfB, void const* ptrPerTokenSfA, void const* ptrPerTokenSfB, + void const* ptrBias, void* dSfC, float const* ptrScaleC, float const* ptrScaleGate, + float const* ptrClampLimit, float const* ptrSwiGluAlpha, float const* ptrSwiGluBeta, + int32_t const* routeMap, float* rowMax, uint32_t* rowMaxBars, + int32_t const* ptrNumNonExitingCtas = nullptr, int32_t const* ptrTotalNumPaddedTokens = nullptr, + int32_t const* ptrCtaIdxXyToBatchIdx = nullptr, int32_t const* ptrCtaIdxXyToMnLimit = nullptr, + int32_t const maxNumCtas = KernelParams::MaxNumCtas) { + static_assert(sizeof(KernelParams) <= 32 * 1024, + "sizeof(KernelParams) has to be less or equal than 32KB"); + + // Create the return struct. + KernelParams params; + + params.ptrRouteMap = routeMap; + params.numTokens = options.mNumTokens; + + params.ptrScaleC = ptrScaleC; + params.ptrScaleGate = ptrScaleGate; + params.ptrClampLimit = ptrClampLimit; + params.ptrSwiGluAlpha = ptrSwiGluAlpha; + params.ptrSwiGluBeta = ptrSwiGluBeta; + + int32_t ctaOffset = 0; + + // Compute totalNumPaddedTokens, ctaIdxXyToBatchIdx and ctaIdxXyToMnLimit if the batch dims are + // known at kernel launch time. Otherwise, these parameters are defined in the device buffers: + // ptrTotalNumPaddedTokens, ptrCtaIdxXyToBatchIdx and ptrCtaIdxXyToMnLimit respectively. + + if (options.mIsStaticBatch) { + params.totalNumPaddedTokens = 0; + for (int b = 0; b < options.mNumBatches; b++) { + int mM = batchM ? options.mBatchedM[b] : options.mM; + int mN = batchM ? options.mN : options.mBatchedN[b]; + + // Skip Tma descriptor creation if expert isn't used + if (mM == 0 || mN == 0) { + continue; + } + + // The number of CTAs. + int32_t numCtas = batchM ? (mM + options.mTileM - 1) / options.mTileM + : (mN + options.mTileN - 1) / options.mTileN; + // The size of the tile. + int32_t tile = batchM ? options.mTileM : options.mTileN; + // The problem size. + int32_t mn = batchM ? mM : mN; + int32_t tokensPerTile = mn; + + // Make sure we do not exceed the launch limit. + if (ctaOffset + numCtas > KernelParams::MaxNumCtas) { + throw std::runtime_error("Too many CTAs"); + } + + for (int32_t cta = 0; cta < numCtas; cta++) { + params.ctaIdxXyToBatchIdx[ctaOffset + cta] = b; + // This is now an identity map and it is no longer needed. + // params.ctaIdxXyToTileIdxMn[ctaOffset + cta] = ctaOffset + cta; + params.ctaIdxXyToMnLimit[ctaOffset + cta] = + std::min((ctaOffset + cta + 1) * tile, ctaOffset * tile + tokensPerTile); + } + ctaOffset += numCtas; + + params.totalNumPaddedTokens += numCtas * tile; + } + } else { + params.ptrTotalNumPaddedTokens = ptrTotalNumPaddedTokens; + params.ptrCtaIdxXyToBatchIdx = ptrCtaIdxXyToBatchIdx; + params.ptrCtaIdxXyToMnLimit = ptrCtaIdxXyToMnLimit; + ctaOffset = maxNumCtas; + } + + if (options.mUseDeepSeekFp8 && options.mDtypeC == tg::Dtype::E4m3) { + params.ptrDqSfsC = reinterpret_cast(dSfC); + } + + params.ptrA = ptrA; + params.ptrB = ptrB; + params.strideInBytesA = options.mK * tg::dtypeGetNumBits(options.mDtypeA) / 8; + params.strideInBytesB = options.mK * tg::dtypeGetNumBits(options.mDtypeB) / 8; + + params.ptrSfA = dSfA; + params.ptrSfB = dSfB; + params.ptrSfC = dSfC; + + if (!batchM) { + // A is the expert + if (0 != options.mM % options.mTileM) { + throw std::runtime_error("0 == mM %% tileM"); + } + params.tileStridePerBatch = options.mM / options.mTileM; + params.nm = options.mM; + // Shape/stride for gmem tensor A. + auto [shapeA, strideA, tileShapeA] = + makeTmaShapeStrideAbc(options, options.mM, options.mN, options.mK, options.mTileM, + options.mTileN, options.mTileK, MatrixType::MatrixA); + // Build tma descriptor for A. + params.tmaA[0] = gemm::buildNdTmaDescriptor(options.mDtypeA, options.mMmaKind, shapeA, strideA, + tileShapeA, const_cast(ptrA)); + + // The input is padded: + // [act0, padding, padding, ... TileN size .., act1, padding, padding, ...] + auto const inputNumTokens = ctaOffset * options.mTileN; + + if (!batchedGemm::doesRouteImplUseLdgsts(options.mRouteImpl)) { + bool useRouteAct = batchedGemm::doesRouteImplUseTma(options.mRouteImpl); + // B is the activation + // Shape/stride for gmem tensor B. + auto [shapeB, strideB, tileShapeB] = makeTmaShapeStrideAbc( + options, options.mM, useRouteAct ? options.mNumTokens : inputNumTokens, options.mK, + options.mTileM, (useRouteAct ? 1 : options.mTileN), options.mTileK, MatrixType::MatrixB); + // Build tma descriptor for B. + params.tmaB[0] = gemm::buildNdTmaDescriptor(options.mDtypeB, options.mMmaKind, shapeB, + strideB, tileShapeB, const_cast(ptrB)); + } + + if (options.mDtypeA == tg::Dtype::E2m1 || options.mDtypeA == tg::Dtype::MxE4m3 || + options.mDtypeA == tg::Dtype::MxE2m1) { + tg::Dtype const dTypeSf = + (options.mDtypeA == tg::Dtype::E2m1) ? tg::Dtype::E4m3 : tg::Dtype::UE8m0; + + // Build TMA descriptor for gmem A block scaling factors. + auto [shapeSfA, strideSfA, tileShapesSfA] = makeTmaShapeStrideSfAb( + options.mM * options.mNumBatches, options.mN, options.mK, MatrixType::MatrixA, + options.mTileM, options.mTileN, options.mTileK, options.mDtypeA, tg::SfLayout::R128c4, + options.mSfReshapeFactor); + params.tmaSfA[0] = gemm::buildSfTmaDescriptor(dTypeSf, shapeSfA, strideSfA, tileShapesSfA, + const_cast(dSfA)); + } + + if (options.mDtypeB == tg::Dtype::E2m1 || options.mDtypeB == tg::Dtype::MxE4m3 || + options.mDtypeB == tg::Dtype::MxE2m1) { + tg::Dtype const dTypeSf = + (options.mDtypeB == tg::Dtype::E2m1) ? tg::Dtype::E4m3 : tg::Dtype::UE8m0; + + if (batchedGemm::doesRouteImplUseTma(options.mRouteImpl)) { + // The input is NOT padded: + // [act0, act1, act2, ...] + + // Build TMA descriptor for gmem B block scaling factors. + int32_t const numEltsPerSf = tg::dtypeNumEltsPerSf(options.mDtypeB); + // Pad number of scaling factors to the nearest multiple of 16 because of the TMA 16B + // alignment requirement. + auto numSfsInK = options.mK / numEltsPerSf; + numSfsInK = ceilDiv(numSfsInK, 16) * 16; + + auto [shapeSfB, strideSfB, tileShapesSfB] = makeTmaShapeStrideAbc( + options, options.mM, options.mNumTokens, numSfsInK, options.mTileM, 1 /* tileN */, + options.mTileK / numEltsPerSf, MatrixType::MatrixB); + params.tmaSfB[0] = gemm::buildNdTmaDescriptor( + dTypeSf, options.mMmaKind, shapeSfB, strideSfB, tileShapesSfB, const_cast(dSfB), + /*doSwizzle*/ true); + } else if (batchedGemm::doesRouteImplUseNoRoute(options.mRouteImpl)) { + // The input is padded: + // [act0, padding, padding, ... TileN size .., act1, padding, padding, ...] + + auto const inputNumTokensSfB = ctaOffset * options.mTileN; + + // Build TMA descriptor for gmem B block scaling factors. + auto [shapeSfB, strideSfB, tileShapesSfB] = + makeTmaShapeStrideSfAb(options.mM, inputNumTokensSfB, options.mK, MatrixType::MatrixB, + options.mTileM, options.mTileN, options.mTileK, options.mDtypeB, + options.mSfLayoutB, options.mSfReshapeFactor); + params.tmaSfB[0] = gemm::buildSfTmaDescriptor(dTypeSf, shapeSfB, strideSfB, tileShapesSfB, + const_cast(dSfB)); + } + } + + // C is the output activation + if (options.mUseTmaStore) { + // Shape/stride for gmem tensor C. + auto [shapeC, strideC, tileShapeC] = makeTmaShapeStrideAbc( + options, options.mM, ctaOffset * options.mTileN, options.mK, options.mTileM, + options.mTileN, options.mTileK, MatrixType::MatrixC); + // Build tma descriptor for C. + params.tmaC[0] = gemm::buildNdTmaDescriptor(options.mDtypeC, tg::MmaKind::Auto, shapeC, + strideC, tileShapeC, ptrC); + } else { + params.ptrC = ptrC; + } + + } else { + // B is the expert + if (0 != options.mN % options.mTileN) { + throw std::runtime_error("0 == mN %% tileN"); + } + params.tileStridePerBatch = options.mN / options.mTileN; + params.nm = options.mN; + // Shape/stride for gmem tensor B. + auto [shapeB, strideB, tileShapeB] = + makeTmaShapeStrideAbc(options, options.mM, options.mN, options.mK, options.mTileM, + options.mTileN, options.mTileK, MatrixType::MatrixB); + // Build tma descriptor for B. + params.tmaB[0] = gemm::buildNdTmaDescriptor(options.mDtypeB, options.mMmaKind, shapeB, strideB, + tileShapeB, const_cast(ptrB)); + + if (options.mRouteImpl == batchedGemm::RouteImpl::NoRoute) { + // A is the activation + // Shape/stride for gmem tensor A. + // The input is padded: + // [act0, padding, padding, ... tileM size .., act1, padding, padding, ...] + auto const inputNumTokens = ctaOffset * options.mTileM; + auto [shapeA, strideA, tileShapeA] = + makeTmaShapeStrideAbc(options, inputNumTokens, options.mN, options.mK, options.mTileM, + options.mTileN, options.mTileK, MatrixType::MatrixA); + // Build tma descriptor for A. + params.tmaA[0] = gemm::buildNdTmaDescriptor(options.mDtypeA, options.mMmaKind, shapeA, + strideA, tileShapeA, const_cast(ptrA)); + } + + if (options.mDtypeA == tg::Dtype::E2m1 || options.mDtypeA == tg::Dtype::MxE4m3 || + options.mDtypeA == tg::Dtype::MxE2m1) { + tg::Dtype const dTypeSf = + (options.mDtypeA == tg::Dtype::E2m1) ? tg::Dtype::E4m3 : tg::Dtype::UE8m0; + + if (options.mRouteImpl == batchedGemm::RouteImpl::NoRoute) { + // The input is padded: + // [act0, padding, padding, ... tileM size .., act1, padding, padding, ...] + auto const inputNumTokensSfA = ctaOffset * options.mTileM; + + // Build TMA descriptor for gmem A block scaling factors. + auto [shapeSfA, strideSfA, tileShapesSfA] = + makeTmaShapeStrideSfAb(inputNumTokensSfA, options.mN, options.mK, MatrixType::MatrixA, + options.mTileM, options.mTileN, options.mTileK, options.mDtypeA, + tg::SfLayout::R128c4, options.mSfReshapeFactor); + params.tmaSfA[0] = gemm::buildSfTmaDescriptor(dTypeSf, shapeSfA, strideSfA, tileShapesSfA, + const_cast(dSfA)); + } + } + + if (options.mDtypeB == tg::Dtype::E2m1 || options.mDtypeB == tg::Dtype::MxE4m3 || + options.mDtypeB == tg::Dtype::MxE2m1) { + tg::Dtype const dTypeSf = + (options.mDtypeB == tg::Dtype::E2m1) ? tg::Dtype::E4m3 : tg::Dtype::UE8m0; + + // Build TMA descriptor for gmem B block scaling factors. + auto [shapeSfB, strideSfB, tileShapesSfB] = makeTmaShapeStrideSfAb( + options.mM, options.mN * options.mNumBatches, options.mK, MatrixType::MatrixB, + options.mTileM, options.mTileN, options.mTileK, options.mDtypeB, options.mSfLayoutB, + options.mSfReshapeFactor); + params.tmaSfB[0] = gemm::buildSfTmaDescriptor(dTypeSf, shapeSfB, strideSfB, tileShapesSfB, + const_cast(dSfB)); + } + + // C is the output activation + if (options.mUseTmaStore) { + // Shape/stride for gmem tensor C. + auto [shapeC, strideC, tileShapeC] = makeTmaShapeStrideAbc( + options, ctaOffset * options.mTileM, options.mN, options.mK, options.mTileM, + options.mTileN, options.mTileK, MatrixType::MatrixC); + // Build tma descriptor for C. + params.tmaC[0] = gemm::buildNdTmaDescriptor(options.mDtypeC, tg::MmaKind::Auto, shapeC, + strideC, tileShapeC, ptrC); + } else { + params.ptrC = ptrC; + } + } + + params.k = options.mK; + params.numBatches = options.mNumBatches; + + params.rank = 0; + params.tpGrpSize = 1; + + params.ptrPartialRowMax = rowMax; + params.ptrRowMaxCompletionBars = rowMaxBars; + + params.ptrNumNonExitingCtas = ptrNumNonExitingCtas; + + // Set the per-token scale factors for MetaFP8 or scale inputs + params.ptrPerTokenSfA = ptrPerTokenSfA; + params.ptrPerTokenSfB = ptrPerTokenSfB; + params.ptrBias = ptrBias; + + return params; +} +#endif +}; // namespace KernelParamsSetup + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace batchedGemm +} // namespace batchedGemm diff --git a/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/KernelParamsDecl.h b/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/KernelParamsDecl.h new file mode 100644 index 0000000000000000000000000000000000000000..63f1d29ce3eb1b83e9e8c096ce47b1385f0ab42a --- /dev/null +++ b/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/KernelParamsDecl.h @@ -0,0 +1,519 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & + * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once +namespace batchedGemm { + +// This is device code + +struct KernelParams { + ////////////////////////////////////////////////////////////////////////////////////////////////// + // + // BatchedGemm parameters. + // + ////////////////////////////////////////////////////////////////////////////////////////////////// + + // Maximum number of CTAs + static constexpr int MaxNumCtas = 2048; + + // NOTE: TMA out-of-bounds optimization for MoE padded tokens: + // + // 1. 2D tensor [hidden, numPaddedTokens] with stride [1, hidden] at pointer p + // 2. Reshape to 3D: [hidden, tileN, numPaddedTokens] with stride [1, hidden, hidden] + // 3. Transform coordinates + // Originally: + // + // Ptr (p ) <-> Coord [:, 0 , ctaIdxY * tileN ] + // + // Equals + // + // Ptr (p - tileN*hidden) <-> Coord [:, tileN , ctaIdxY * tileN ] + // + // Equals + // + // Ptr (p - tileN*hidden) <-> Coord [:, tileN - batchEnd, ctaIdxY * tileN + batchEnd] + // + // In effect, we load exactly `batchEnd` tokens from same global address, reducing wasted traffic. + // + // Due to the restrictions of TMA unit, the above operations requires the TMA descriptor and the + // underlying buffer be constructed differently: + // - Requires valid buffer at (p - tileN * hidden) - needs prepending `tileN` tokens. + // - TMA outermost dimension must be extended by `tileN` or loads will OOB in the rightmost side. + + // TMA descriptor for A. + // Must be setup using gemm::buildNdTmaDescriptor with shapes and strides from + // makeTmaShapeStrideAbc. + // + // If batchM: + // Logical shape is [sum(divUpMul(M[bi], tileM) for bi in B), K]. + // Logical strides are [K, 1]. + // Tile box shape is [tileM, tileK]. + // Tile box strides are [tileK, 1]. + // + // If batchN: + // If layoutA is MatrixLayout::MajorK + // Logical shape is [B, divUpMul(M, tileM), K]. + // Logical strides are [divUpMul(M, tileM) * K, K, 1]. + // Tile box shape is [1, tileM, tileK]. + // Tile box strides are [0, tileK, 1]. + // If layoutA is MatrixLayout::Mn + // Logical shape is [B, K, divUpMul(M, tileM)]. + // Logical strides are [K * divUpMul(M, tileM), divUpMul(M, tileM), 1]. + // Tile box shape is [1, tileK, tileM]. + // Tile box strides are [0, tileM, 1]. + // If layoutA is MatrixLayout::BlockMajorK + // Logical shape is [B, K / blockK, divUpMul(M, tileM), blockK]. + // Logical strides are [K * divUpMul(M, tileM), divUpMul(M, tileM) * blockK, blockK, 1]. + // Tile box shape is [1, tileK / min(blockK, tileK), tileM, min(blockK, tileK)]. + // Tile box strides are [0, tileM * min(blockK, tileK), min(blockK, tileK), 1]. + // where blockK is 128B. + // + // Dtype is set from options.mDtypeA. + CUtensorMap tmaA[1]; + + // TMA descriptor for B. + // Must be setup using gemm::buildNdTmaDescriptor with shapes and strides from + // makeTmaShapeStrideAbc. + // + // If batchM: + // If layoutB is MatrixLayout::MajorK + // Logical shape is [B, divUpMul(N, tileN), K]. + // Logical strides are [divUpMul(N, tileN) * K, K, 1]. + // Tile box shape is [1, tileN, tileK]. + // Tile box strides are [0, tileK, 1]. + // If layoutB is MatrixLayout::MajorMn + // Logical shape is [B, K, divUpMul(N, tileN)]. + // Logical strides are [K * divUpMul(N, tileN), divUpMul(N, tileN), 1]. + // Tile box shape is [1, tileK, tileN]. + // Tile box strides are [0, tileN, 1]. + // If layoutB is MatrixLayout::BlockMajorK + // Logical shape is [B, K / blockK, divUpMul(N, tileN), blockK]. + // Logical strides are [K * divUpMul(N, tileN), divUpMul(N, tileN) * blockK, blockK, 1]. + // Tile box shape is [1, tileK / min(blockK, tileK), tileN, min(blockK, tileK)]. + // Tile box strides are [0, tileN * min(blockK, tileK), min(blockK, tileK), 1]. + // where blockK is 128B. + // + // If batchN: + // Logical shape is [sum(divUpMul(N[bi], tileN) for bi in B), K]. + // Logical strides are [K, 1]. + // Tile box shape is [tileN, tileK]. + // Tile box strides are [tileK, 1]. + // + // Dtype is set from options.mDtypeB. + CUtensorMap tmaB[1]; + + // TMA descriptor for C, (when useTmaStore is true) + // Must be setup using gemm::buildNdTmaDescriptor with shapes and strides from + // makeTmaShapeStrideAbc. + // + // If batchM: + // Logical shape is [sum(divUpMul(M[bi], tileM) for bi in B), N]. + // Logical strides are [N, 1]. + // Tile box shape is [epilogueTileM, epilogueTileN]. + // Tile box strides are [epilogueTileN, 1]. + // + // If batchN: + // Logical shape is [sum(divUpMul(N[bi], tileN) for bi in B), M]. + // Logical strides are [M, 1]. + // Tile box shape is [epilogueTileN, epilogueTileM]. + // Tile box strides are [epilogueTileM, 1]. + // + // Dtype is set from options.mDtypeC. + CUtensorMap tmaC[1]; + + // TMA descriptor for the block scaling factors for A, for MxFp{4,8} and NvFp4 formats. + // Must be setup using gemm::buildSfTmaDescriptor with shapes and strides from + // makeTmaShapeStrideSfAb. + // The layout of scaling factors for A is always R128c4. + // + // Let P be the number of elements per SF. P=16 for NvFp4, P=32 for Mx formats. + // M must be a multiple of 128. + // K must be a multiple of 4P. + // The "logical" shape is: [paddedM, K / P], where paddedM is + // sum(divUpMul(M[bi], tileM) for bi in B) if batchM, + // otherwise divUpMul(M, TileM) * B. + // The R128c4 layout is: [paddedM / 128, K / P / 4, 512]. + // The shape we use for TMA is: [paddedM / 128, K / P / 4, 2, 256]. + // + // Dtype is Dtype::E4m3 for NvFp4, Dtype::UE8m0 for Mx formats. + CUtensorMap tmaSfA[1]; + + // TMA descriptor for the block scaling factors for B, for MxFp{4,8} and NvFp4 formats. + // Must be setup using gemm::buildSfTmaDescriptor with shapes and strides from + // makeTmaShapeStrideSfAb. + // The layout of block scaling factors for B is controlled by options.mSfLayoutB. + // + // Let P be the number of elements per SF. P=16 for NvFp4, P=32 for Mx formats. + // The "logical" shape is: [paddedN, K / 16] + // where paddedN is sum(divUpMul(N[bi], tileN) for bi in B) if batchN, + // otherwise divUpMul(N, TileN) * B. + // + // If the layout is R128c4, + // paddedN must be a multiple of 128. + // K must be a multiple of 4P. + // The R128c4 layout is: [paddedN / 128, K / P / 4, 512] + // The shape we use for TMA is: [paddedN / 128, K / P / 4, 2, 256] + // + // If the layout is R8c4, + // paddedN must be a multiple of 8. + // K must be a multiple of 4P. + // The R8c4 layout is: [paddedN / 8, K / P / 4, 32] + // The shape we use for TMA is: [paddedN / 8, K / P / 4 / repeats, repeats * 32] + // where repeats = min(tileK / P / 4, 8) + // + // Dtype is Dtype::E4m3 for NvFp4, Dtype::UE8m0 for Mx formats. + CUtensorMap tmaSfB[1]; + + // The input matrix A. + // If (routeAct == true && batchM), the shape is [M, K]. tmaA is not used. + // Otherwise, check layout of tmaA to see the shape and strides. + void const* ptrA{nullptr}; + + // The stride for matrix A in bytes. + // Equals to K * dtypeGetNumBits(dtypeA) / 8. + uint64_t strideInBytesA; + + // The input matrix B. + // If (routeAct == true && batchN), the shape is [N, K]. tmaB is not used. + // Otherwise, check layout of tmaB to see the shape and strides. + void const* ptrB{nullptr}; + // The stride for matrix B in bytes. + // Equals to K * dtypeGetNumBits(dtypeB) / 8. + uint64_t strideInBytesB; + + // The output matrix C. Check "logical" layout of tmaC to see the shape and strides. + void* ptrC{nullptr}; + + // Inputs and output are MxFp{4,8}, Fp8, NvFp4. + // The scaling factors to apply to the output - can be used to incorporate input scaling factors + // as described below: C = SEncC * act(SDecA * SDecB * A * Bl) . (SDecA * SDecB * A * Br) + // -> ScaleGate = SDecA * SDecB + // ScaleC = SDecA * SDecB * SEncC + // + // Only the inputs are MxFp{4,8}, Fp8, NvFp4. + // C = act(SDecA * SDecB * A * Bl) . (SDecA * SDecB * A * Br) + // -> ScaleGate = SDecA * SDecB + // ScaleC = SDecA * SDecB + // + // Only the output is MxFp{4,8}, Fp8, NvFp4. + // C = SEncC * act(A * Bl) . (A * Br) + // -> ScaleGate = 1 + // ScaleC = SEncC + // + // The output tensor scaling factor for MxFp{4,8}, Fp8, NvFp4 and DeepSeek FP8 quantization. + // TensorRT-LLM API requires a scaling factor on the device. + // Shape is [B]. One scaling factor per tensor in batch. + float const* ptrScaleC{nullptr}; + + // The output gate scale for MxFp{4,8}, Fp8, NvFp4 and DeepSeek FP8 quantization. + // TensorRT-LLM API requires a scaling factor on the device. + // Shape is [B]. One scaling factor per tensor in batch. + float const* ptrScaleGate{nullptr}; + + // The clamp limit before the activation. + // Shape is [B]. + // Clamp is INF if nullptr. + // If applied on SwiGlu, it will be: + // + // x_glu = x_glu.clamp(min=None, max=limit) + // x_linear = x_linear.clamp(min=-limit, max=limit) + float const* ptrClampLimit{nullptr}; + + // The alpha and beta for SwiGlu. + // Shape is [B]. One alpha and one beta per tensor in batch. + // Alpha is 1.f if nullptr. + // Beta is 0.f if nullptr. + // The formula: + // + // out_glu = x_glu * torch.sigmoid(alpha * x_glu) * (x_linear + beta) + float const* ptrSwiGluAlpha{nullptr}; + float const* ptrSwiGluBeta{nullptr}; + + // The K dimension. It is the hidden dimension of the input matrices. + int32_t k; + + // The non-batched dimension. + // It is N if batchM, otherwise M. + int32_t nm; + + // Tile stride per batch for the non-batched dimension. + // It is N / TileN if batchM, otherwise M / TileM. + int32_t tileStridePerBatch; + + // TODO get rid of that. + // DeepSeek FP8 scaling factors for C + float* ptrDqSfsC{nullptr}; + + // The block scaling factors for A. + // The pointer must always be set regardless of the quantization recipe. + // If (routeAct == true && batchM), the shape is [M, K / 16]. tmaSfA is not used. + // For the layout (r128c4), see below. + // Otherwise, + // If MxFp{4,8} and NvFp4 formats are used, + // check the "logical" layout of tmaSfA to see the shape and strides. + // The dtype is Dtype::E4m3. + // + // If DeepSeek FP8 quantization recipe is used, + // If batchM: + // The shape is [K / 128, paddedM], + // where paddedM is sum(divUpMul(M[bi], tileM) for bi in B). + // If batchN: + // The shape is [M / 128, K / 128], + // The rightmost dimension is contiguous in memory. + // The dtype is Dtype::Float32. + void const* ptrSfA{nullptr}; + + // The block scaling factors for B. + // The pointer must always be set regardless of the quantization recipe. + // If (routeAct == true && batchN), the shape is [N, K / 16]. tmaSfB is not used. + // For the layout (r128c4, r8c4), see below. + // Otherwise, + // If MxFp{4,8} and NvFp4 formats are used, + // check the layout of tmaSfB to see the shape and strides. + // The dtype is Dtype::E4m3. + // + // If DeepSeek FP8 quantization recipe is used, + // If batchM: + // The shape is [N / 128, K / 128], + // If batchN: + // The shape is [K / 128, paddedN], + // where paddedN is sum(divUpMul(N[bi], tileN) for bi in B). + // The rightmost dimension is contiguous in memory. + // The dtype is Dtype::Float32. + void const* ptrSfB{nullptr}; + + // The per-token scaling factors from scale A. + // + // This is used for either: + // * Per-token scaling factor quantization schemes, such as MetaFP8. The dtype is Dtype::Float32 + // * When the routing scales are applied to the input activations (only when output is not + // transposed). The dtype is Dtype::Bfloat16 + // + // if (batchM (A is activations)): + // Logical shape is [sum(divUpMul(M[bi], tileM) for bi in B)] + // + // if (batchN (A is weights)): + // Logical shape is [B, divUpMul(M, tileM)] + // + void const* ptrPerTokenSfA{nullptr}; + + // The per-token scaling factors from scale B. + // + // This is used for either: + // * Per-token scaling factor quantization schemes, such as MetaFP8. The dtype is Dtype::Float32 + // * When the routing scales are applied to the input activations (only when output is + // transposed). The dtype is Dtype::Bfloat16 + // + // if (batchM (B is weights)): + // Logical shape is [B, divUpMul(N, tileN)] + // + // if (batchN (B is activations)): + // Logical shape is [sum(divUpMul(N[bi], tileN) for bi in B)] + void const* ptrPerTokenSfB{nullptr}; + + // The bias applied after the GEMM and before the activation function. + // The bias is applied before applying the global scaling factor. I.e. + // C = act(A * B + bias') * scaleC + // scaleC = dequantA * dequantB * quantC + // Thus, the bias' = bias / (dequantA * dequantB), where the bias is the original bias. + // + // If batchM, BiasType must be N, and bias shape is [B, N]. + // The bias is broadcasted along the M dimension. + // + // If batchNm BiasType must be M, and bias shape is [B, M]. + // The bias is broadcasted along the N dimension. + // + // The dtype is float32. + void const* ptrBias{nullptr}; + + // The output block scaling factors for C. + // + // If MxFp{4,8} and NvFp4 formats are used, + // The "logical" shape is: + // if batchM: [paddedM, N / 16] + // if batchN: [paddedN, M / 16] + // where paddedM is sum(divUpMul(M[bi], tileM) for bi in B), + // where paddedN is sum(divUpMul(N[bi], tileN) for bi in B). + // + // If the layout is R128c4, + // paddedOuter must be a multiple of 128. + // inner must be a multiple of 64. + // The R128c4 layout is: [paddedOuter / 128, inner / 16 / 4, 512] + // The shape we use for TMA is: [paddedOuter / 128, inner / 16 / 4, 2, 256] + // where inner = N if batchM, otherwise M. + // where paddedOuter = paddedM if batchM, otherwise paddedN. + // + // If the layout is R8c4, + // paddedOuter must be a multiple of 8. + // inner must be a multiple of 64. + // The R8c4 layout is: [paddedOuter / 8, inner / 16 / 4, 32] + // The shape we use for TMA is: [paddedOuter / 8, inner / 16 / 4 / repeats, repeats * 32] + // where repeats = min(tileInner / 16 / 4, 8), + // where tileInner = tileN if batchM, otherwise tileM, + // where paddedOuter = paddedM if batchM, otherwise paddedN. + // where inner = N if batchM, otherwise M. + // + // The dtype is Dtype::E4m3. + // + // If DeepSeek FP8 quantization recipe is used, + // If batchM: + // The shape is [N / 128, paddedM], + // where paddedM is sum(divUpMul(M[bi], tileM) for bi in B). + // If batchN: + // The shape is [M / 128, paddedN], + // where paddedN is sum(divUpMul(N[bi], tileN) for bi in B). + // The rightmost dimension is contiguous in memory. + // The dtype is Dtype::Float32. + void* ptrSfC{nullptr}; + + ////////////////////////////////////////////////////////////////////////////////////////////////// + // + // Routing activations parameters. + // + ////////////////////////////////////////////////////////////////////////////////////////////////// + // These params are used when the kernel is configured with -routeAct true. + // The inputs are not padded, but the outputs are padded to divUpMul(M[bi], tileM) for batchM or + // divUpMul(N[bi], tileN) for batchN. + // If -routeAct is false, the params are not used and should be set to zero. + + // The routeMap for the input tokens. + // Map of expanded token index (counting the previous padded tokens) to the batch index + // the token belongs to. + // The shape is + // [sum(divUpMul(M[bi], tileM) for bi in B)] for batchM + // [sum(divUpMul(N[bi], tileN) for bi in B)] for batchN + // The dtype is int32_t. + // + // There are 3 tokens [0, 1, 2] such that [0, 1] belong to batch [B0] and [2] to batch [B1]. + // Let's assume that the padded size is 4. + // + // The expanded indices for tokens [0, 1, 2] are: + // expandedIdx[0] = 0 + // expandedIdx[1] = 1 + // expandedIdx[2] = divUpMul(2, 4) + 0 = 4 + // + // The route map is [B0, B0, X, X, B1, X, X, X] where X could be any value. + int32_t const* ptrRouteMap{nullptr}; + + // Total number of unpadded inputs + int32_t numTokens; + + // Total number of batches + int32_t numBatches; + + ////////////////////////////////////////////////////////////////////////////////////////////////// + // + // Batching information parameters. + // + ////////////////////////////////////////////////////////////////////////////////////////////////// + + // In some cases, some CTAs must early-exit. E.g. when the grid size is set statically, but the + // actual workload is decided at runtime. This element on the device contains the number of CTAs + // that do not early-exit. The number corresponds to the X dim of the grid when the output is not + // transposed (i.e. batchM). To the Y dim, otherwise. + // The size is 1 and the dtype is int32_t. + // Used if isStaticBatch == false, otherwise set to nullptr. + // The pointer points to a scalar and the dtype is int32_t. The pointed value must be >= 0. + int32_t const* ptrNumNonExitingCtas{nullptr}; + + // Pointer to total number of padded tokens. + // Computed as + // int32_t totalNumPaddedTokens{0}; + // for (int bi = 0; bi < options.mNumBatches; bi++) { + // totalNumPaddedTokens += batchM ? divUpMul(options.mBatchedM[bi], options.mTileM) + // : divUpMul(options.mBatchedN[bi], options.mTileN); + // } + // The size is 1 and the dtype is int32_t. + // If isStaticBatch == true, ptrTotalNumPaddedTokens should be set to nullptr and + // totalNumPaddedTokens is used. + int32_t const* ptrTotalNumPaddedTokens{nullptr}; + + // Pointer to the map from the CTA index (in X/Y dim) to the batch index. + // Maps CTA index in batch dim (i.e. blockDim.x if batchM, otherwise blockDim.y) + // to batch index. + // E.g. with listM = 128,255,32 and tileM = 128, should be equal to + // ctaIdxXyToBatchIdx = [0, 1, 1, 2] + // If isStaticBatch == true, ptrCtaIdxXyToBatchIdx should be set to nullptr and ctaIdxXyToBatchIdx + // is used. + int32_t const* ptrCtaIdxXyToBatchIdx{nullptr}; + + // Pointer from the CTA index X/Y to the expanded tile index where the expanded tile index is + // computed as: + // + // int expandedIdx = 0; + // for (int bi = 0; bi < batchIdx-1; ++bi) { + // expandIdx = divUpMul(numTokens[bi], TileM/N); + // } + // expandIdx += + // E.g. with numTokens = [128,255,32] and tileM = 128, should be equal to + // ptrCtaIdxXyToMnLimit = [128, 256, 383, 416] + int32_t const* ptrCtaIdxXyToMnLimit{nullptr}; + + // Total number of padded tokens - used as the stride for the activation and C scaling factors. + // Check ptrTotalNumPaddedTokens to see how it is computed. + // If isStaticBatch == true, totalNumPaddedTokens is used, otherwise ptrTotalNumPaddedTokens. + int32_t totalNumPaddedTokens; + + // A map from CTA index X/Y to batch index. + // Check ptrCtaIdxXyToBatchIdx to see how it is computed. + // If isStaticBatch == true, ctaIdxXyToBatchIdx is used, otherwise ptrCtaIdxXyToBatchIdx. + int32_t ctaIdxXyToBatchIdx[MaxNumCtas]; + + // **Expanded** limits for the batched dimension: + // tile * ctaIdxXyToTileIdxMn[ctaIdxXy] -> ctaIdxXyToMnLimit[ctaIdxXy] + // Check ptrCtaIdxXyToMnLimit to see how it is computed. + // If isStaticBatch == true, ctaIdxXyToMnLimit is used, otherwise ptrCtaIdxXyToMnLimit. + int32_t ctaIdxXyToMnLimit[MaxNumCtas]; + + ////////////////////////////////////////////////////////////////////////////////////////////////// + // + // All-reduce parameters. + // + ////////////////////////////////////////////////////////////////////////////////////////////////// + + // The rank id of the current device in the multi-gpu space. + int rank; + // The number of peer devices in tensor-parallel group. + int tpGrpSize; + + ////////////////////////////////////////////////////////////////////////////////////////////////// + // + // GatedAct parameters. + // + ////////////////////////////////////////////////////////////////////////////////////////////////// + + // Pointer for partial row max for DeepSeek FP8 recipe. + // This is temporary storage for the row max results. + // If batchM, the shape is [2, totalNumPaddedTokens, N / 128] and the dtype is float. + // Otherwise, the shape is [2, totalNumPaddedTokens, M / 128] and the dtype is float. + float* ptrPartialRowMax{nullptr}; + + // Flags in global memory that sync on "exit" for row max computation. + // The shape is [numTilesM * numTilesN / 2] and the dtype is uint32_t, where + // if batchM, + // numTilesM = divUp(totalNumPaddedTokens, tileM). + // numTilesN = divUp(N, tileN). + // Otherwise, + // numTilesM = divUp(M, tileM). + // numTilesN = divUp(totalNumPaddedTokens, tileN). + // + // The memory must be set to 0 before the kernel launch. + uint32_t* ptrRowMaxCompletionBars{nullptr}; +}; + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace batchedGemm diff --git a/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/KernelTraits.h b/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/KernelTraits.h new file mode 100644 index 0000000000000000000000000000000000000000..b13ad08c8a2f2097ac7121a172682213a709d0d7 --- /dev/null +++ b/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/KernelTraits.h @@ -0,0 +1,638 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & + * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include +#include + +#include "Enums.h" +#include "trtllm/gen/CommonUtils.h" +#include "trtllm/gen/DtypeDecl.h" + +namespace batchedGemm { + +namespace gemm { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace tg = trtllm::gen; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// Structure to manage memory allocation with configurable reuse +class MemAllocatorHelper { + public: + // The default constructor. + MemAllocatorHelper() {} + + // Constructor to initialize chunk sizes, alignments, and reuse flags + MemAllocatorHelper(std::vector> const& sizes, + std::vector const& reuse, std::vector const& names) + : mNumBytesAndAlignmentPerSmemChunk(sizes), mFirstChunkReuse(reuse), mSmemChunkNames(names) {} + + // Function to calculate the size of the array from 0 to jj chunks + int32_t getOffsetBeforeChunk(int jj) const { + int32_t totalSize = 0; + for (int32_t ii = 0; ii < jj; ++ii) { + auto const& elem = mNumBytesAndAlignmentPerSmemChunk[ii]; + auto paddedSize = getSizePaddedToAlignment(elem.first, elem.second); + // If SMEM chunk is reused but the size of the current chunk is + // larger than currently counted size + if (mFirstChunkReuse[ii] && paddedSize > totalSize) { + // Set new size to the size of the current chunk. + // E.g. possible in case of + // mNumBytesAndAlignmentPerSmemChunk = {{1, 1}, {1, 1}, {1024, 1}} + // mFirstChunkReuse = {false, false, true} + // The last chunk is larger than the first plus second, so total size is 1024. + totalSize = paddedSize; + } else if (!mFirstChunkReuse[ii]) { + totalSize += paddedSize; + } + } + return totalSize; + } + + // Returns the offset of the ith chunk + int32_t getChunkOffsetByName(std::string const& name) const { + for (size_t ii = 0; ii < mSmemChunkNames.size(); ++ii) { + if (mSmemChunkNames[ii] == name) { + return getChunkOffset(ii); + } + } + throw std::runtime_error("Name not found: " + name); + } + + // Returns the first chunk reuse flag given chunk name. + int getFirstChunkReuseFlagByName(std::string const& name) const { + for (size_t ii = 0; ii < mSmemChunkNames.size(); ++ii) { + if (mSmemChunkNames[ii] == name) { + return getFirstChunkReuseFlag(ii); + } + } + throw std::runtime_error("Name not found: " + name); + } + + // Function to calculate the total size of the SMEM array + int32_t getTotalSize() const { + return getOffsetBeforeChunk(static_cast(mNumBytesAndAlignmentPerSmemChunk.size())); + } + + private: + int32_t getChunkOffset(int32_t ii) const { + if (mFirstChunkReuse[ii]) { + // Reuse the offset of the 0th chunk. + return getChunkOffset(0); + } + + // Get offset of ii chunks. + auto offset = getOffsetBeforeChunk(ii); + // Ensure alignment for the current chunk + return getSizePaddedToAlignment(offset, mNumBytesAndAlignmentPerSmemChunk[ii].second); + } + + // Returns the first chunk reuse flag for the ith chunk. + int getFirstChunkReuseFlag(int32_t ii) const { return mFirstChunkReuse[ii]; } + + // Print the contents of this object. + void print() const { + for (size_t ii = 0; ii < mNumBytesAndAlignmentPerSmemChunk.size(); ++ii) { + printf("Chunk %zd %s: %d bytes, %d alignment, reuse %s, offset %d\n", ii, + mSmemChunkNames[ii].c_str(), mNumBytesAndAlignmentPerSmemChunk[ii].first, + mNumBytesAndAlignmentPerSmemChunk[ii].second, mFirstChunkReuse[ii] ? "true" : "false", + getChunkOffset(ii)); + } + } + + private: + // Helper function to calculate padded size + int32_t getSizePaddedToAlignment(int32_t size, int32_t alignment) const { + assert((alignment & (alignment - 1)) == 0); + return (size + alignment - 1) & ~(alignment - 1); + } + + private: + // Sizes and alignment requirements of each chunk + // NOTE: be careful and make sure that the memory dependency is clear and + // chunks in the beginning of the SMEM can be overwritten. + std::vector> mNumBytesAndAlignmentPerSmemChunk; + // Chunk reuse configuration. True at ith position means that ith chunk starts at smemOffset = 0. + std::vector mFirstChunkReuse; + // Buffer names for inspection purposes. + std::vector mSmemChunkNames; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline int getNumSmemBitsPerElt(tg::Dtype dtype, tg::MmaKind mmaKind) { + if (mmaKind == tg::MmaKind::Auto) { + throw std::runtime_error("mmaKind != tg::MmaKind::Auto"); + } + if (mmaKind == tg::MmaKind::MxFp8Fp6Fp4) { + return 8; + } else { + return tg::dtypeGetNumBits(dtype); + } +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +class KernelTraits { + public: + // The default constructor. + KernelTraits() {} + + // The constructor. + KernelTraits(tg::Dtype dtypeA, tg::Dtype dtypeB, tg::Dtype dtypeC, tg::Dtype dtypeAcc, + tg::Dtype dtypeMmaA, tg::Dtype dtypeMmaB, tg::MmaKind mmaKind, int32_t tileM, + int32_t tileN, int32_t tileK, int32_t epilogueTileM, int32_t epilogueTileN, + int32_t numStages, int32_t numStagesMma, int32_t numSlicesForSplitK, + int32_t numSlicesForSliceK, SplitK splitK, bool useTmaStore, bool transposeMmaOutput, + AllReduceAlgo allReduceAlgo, bool usePersistentScheduler, bool useDeepSeekFp8, + bool usePerTokenSfA, bool usePerTokenSfB, BiasType biasType) + : mMmaKind{mmaKind} { + // + // SMEM + // + { + // [smemA ] (1024B aligned) + // [smemB ] (1024B aligned) + // [smemBShuffle ] (1024B aligned) + // [gmemC0 ] (1024B aligned) (if needed) + // [gmemC1 ] (1024B aligned) (if needed) + // [rowMax ] (16B aligned) (if needed) + // [sliceK ] (16B aligned) (if needed) + // [per-token SF ] (16B aligned) (if needed) + // [bias ] (16B aligned) (if needed) + // + // SMEM for smemA and smemB might be repurposed and used for gmemC0 and gmemC1: + // + // [..smemA..][..smemB..][..smemBShuffle..] + // [..gmemC0..][..gmemC1..][..rowMax..][..sliceK..][..per-token SF..][..bias..] + // + + if (mMmaKind == tg::MmaKind::Auto) { + mMmaKind = dtypeGetMmaKind(dtypeMmaA, dtypeMmaB); + } + + std::vector> numBytesAndAlignmentPerSmemChunk; + std::vector firstChunkReuseSmem; + // Buffer names for inspection purposes. + std::vector smemChunkNames; + + // LoadA + { + // Number of bytes in load A shared memory. + auto const numSmemBytesLoadA = + numStages * tileM * tileK * getNumSmemBitsPerElt(dtypeA, mMmaKind) / 8 /* bits */; + // Number of bytes for load A alignment for TMA load. + auto const numBytesAlignmentLoadA = 1024; + // loadA is already at first chunk. No need to reuse it. + auto const reuseChunksSmemLoadA = false; + // Add info. + smemChunkNames.emplace_back("smemLoadA"); + numBytesAndAlignmentPerSmemChunk.emplace_back( + std::make_pair(numSmemBytesLoadA, numBytesAlignmentLoadA)); + firstChunkReuseSmem.emplace_back(reuseChunksSmemLoadA); + } + + // LoadB + { + // Number of bytes in load B shared memory. + auto const numSmemBytesLoadB = + numStages * tileN * tileK * getNumSmemBitsPerElt(dtypeB, mMmaKind) / 8 /* bits */; + // Number of bytes for load B alignment for TMA load. + auto const numBytesAlignmentLoadB = 1024; + // No need to reuse the first chunk. + auto const reuseChunksSmemLoadB = false; + // Add info. + smemChunkNames.emplace_back("smemLoadB"); + numBytesAndAlignmentPerSmemChunk.emplace_back( + std::make_pair(numSmemBytesLoadB, numBytesAlignmentLoadB)); + firstChunkReuseSmem.emplace_back(reuseChunksSmemLoadB); + } + + // SmemBShuffle + // FIXME: we should be able either: + // - Do modification in-place. For that we need to resolve pipeline dependency between + // smemB -> shuffleSmemB -> mma + // - Do 4 TMA SW32 loads or several LDGSTS loads. + { + // Number of bytes in save shuffled B in shared memory. + auto const numSmemBytesLoadB = + numSlicesForSliceK > 1 + ? numStages * tileN * tileK * getNumSmemBitsPerElt(dtypeB, mMmaKind) / 8 /* bits */ + : 0; + // Number of bytes for load B alignment for TMA load. + auto const numBytesAlignmentLoadB = 1024; + // No need to reuse the first chunk. + auto const reuseChunksSmemLoadB = false; + + // Add info. + smemChunkNames.emplace_back("smemBShuffle"); + numBytesAndAlignmentPerSmemChunk.emplace_back( + std::make_pair(numSmemBytesLoadB, numBytesAlignmentLoadB)); + firstChunkReuseSmem.emplace_back(reuseChunksSmemLoadB); + } + + // GmemC + // FIXME we might need to fix this for GemmGatedAct, it needs less SMEM to store gated output. + for (int resIdx = 0; resIdx < 2; ++resIdx) { + // Type of the data in the SMEM for GmemC + auto dtypeSmemC = dtypeC; + if (allReduceAlgo == AllReduceAlgo::TwoShot || numSlicesForSplitK > 1) { + dtypeSmemC = dtypeAcc; + } + // Smem is used for GmemC output tile for TMA store and SplitK in CGA. + bool usesSmemForGmemC = useTmaStore || doesSplitKUseDsmem(splitK); + // SMEM for at leader CTA in DSMEM split-k contains K slices. + auto extraGmemCMultiplier = doesSplitKUseDsmem(splitK) ? numSlicesForSplitK : 1; + if (numSlicesForSliceK > 1) { + // TileN is expanded in N dimension for slice-K. + extraGmemCMultiplier *= numSlicesForSliceK; + } + + if (resIdx != 0 && !useDeepSeekFp8) { + // No data for Epilogue1 in case of non-DeepSeek GEMM. + extraGmemCMultiplier = 0; + } + + // Number of bytes to store the output in smem. + auto const numBytesSmemStoreC = usesSmemForGmemC + ? extraGmemCMultiplier * epilogueTileM * epilogueTileN * + tg::dtypeGetNumBits(dtypeSmemC) / 8 /* bits */ + : 0; + // Number of bytes for store C alignment for TMA store. + auto const numBytesAlignmentStoreC = 1024; + // gmemC reuses loadAb memory for split-K in DSMEM. + // Epilogue1 does not reuse and continues after the memory allocated Epilogue0 + // NOTE: we can always reuse loadAb SMEM as long as we don't have persistent scheduler. + auto const reuseFirstChunksSmemStoreC = + doesSplitKUseDsmem(splitK) && resIdx == 0 && !usePersistentScheduler; + + // Add info. + smemChunkNames.emplace_back("smemGmemC" + std::to_string(resIdx)); + numBytesAndAlignmentPerSmemChunk.emplace_back( + std::make_pair(numBytesSmemStoreC, numBytesAlignmentStoreC)); + firstChunkReuseSmem.emplace_back(reuseFirstChunksSmemStoreC); + } + + // RowMax + { + // Number of dqSfsC per CTA. + auto const numDqSfsCPerCta = transposeMmaOutput ? tileM : tileN; + // Number of bytes for rowMax in SMEM. + auto const numBytesSmemRowMax = (useDeepSeekFp8 ? numDqSfsCPerCta : 0) * + tg::dtypeGetNumBits(tg::Dtype::Fp32) / 8 /* bits */; + // Number of bytes alignment for rowMax in SMEM. + auto const numBytesAlignmentRowMax = 16; + + // Add info. + smemChunkNames.emplace_back("smemRowMax"); + numBytesAndAlignmentPerSmemChunk.emplace_back( + std::make_pair(numBytesSmemRowMax, numBytesAlignmentRowMax)); + firstChunkReuseSmem.emplace_back(false); + } + + // SliceK + { + // Real tile size before slice-K reduction. + auto const tileSize = + numSlicesForSliceK > 1 ? numSlicesForSliceK * tileM * numSlicesForSliceK * tileN : 0; + // Number of bytes for tile in SMEM. + auto const numBytesSmemTile = tileSize * tg::dtypeGetNumBits(dtypeAcc) / 8 /* bits */; + // Number of bytes alignment for rowMax in SMEM. + auto const numBytesAlignmentTile = 16; + + // Add info. + smemChunkNames.emplace_back("smemSliceK"); + numBytesAndAlignmentPerSmemChunk.emplace_back( + std::make_pair(numBytesSmemTile, numBytesAlignmentTile)); + firstChunkReuseSmem.emplace_back(false); + } + + // Per-token Scale Factors + { + // Number of bytes for per-token scale factors + auto const numBytesSmemPerTokenSf = (usePerTokenSfA ? (tileM) * sizeof(float) : 0) + + (usePerTokenSfB ? (tileN) * sizeof(float) : 0); + // Number of bytes alignment for per-token scale factors + auto const numBytesAlignmentPerTokenSf = 16; + // Add info. + smemChunkNames.emplace_back("smemPerTokenSf"); + numBytesAndAlignmentPerSmemChunk.emplace_back( + std::make_pair(numBytesSmemPerTokenSf, numBytesAlignmentPerTokenSf)); + firstChunkReuseSmem.emplace_back(false); + } + + // Bias + { + int32_t numBytesSmemBias = 0; + if (isBiasTypeN(biasType)) { + numBytesSmemBias = tileN * sizeof(float); + } else if (isBiasTypeM(biasType)) { + numBytesSmemBias = tileM * sizeof(float); + } else if (isBiasTypeMn(biasType)) { + numBytesSmemBias = tileM * tileN * sizeof(float); + } + // Number of bytes alignment for bias + auto const numBytesAlignmentBias = 16; + // Add info. + smemChunkNames.emplace_back("smemBias"); + numBytesAndAlignmentPerSmemChunk.emplace_back( + std::make_pair(numBytesSmemBias, numBytesAlignmentBias)); + firstChunkReuseSmem.emplace_back(false); + } + + // Per-block absolute maximum for multi-warp reduction. + { + // Number of bytes: number of epilogue warps * number of tile columns. + // TODO: avoid allocating this memory when it's not needed (it's only for MxFp8 + fusedAct) + auto const numBytesSmemBlockAmax = transposeMmaOutput ? 4 * tileN * sizeof(float) : 0; + // Number of bytes alignment. + auto const numBytesAlignmentBlockAmax = 16; + // Add info. + smemChunkNames.emplace_back("smemBlockAmax"); + numBytesAndAlignmentPerSmemChunk.emplace_back( + std::make_pair(numBytesSmemBlockAmax, numBytesAlignmentBlockAmax)); + firstChunkReuseSmem.emplace_back(false); + } + + // SmemConstSfBuf + // A buffer used to copy constant values to TMEM. + { + // Do we need the buffer? + bool const useConstSfBuf = dtypeB == tg::Dtype::E4m3 && dtypeMmaB == tg::Dtype::MxE4m3; + // Number of bytes for the buffer. + auto const numSmemBytesConstSfBuf = useConstSfBuf ? 512 : 0; + // Number of bytes for the alignment of the buffer. + auto const numBytesAlignmentConstSfBuf = 16; + // No need to reuse the first chunk. + auto const reuseChunksSmemConstSfBuf = false; + + // Add info. + smemChunkNames.emplace_back("smemConstSfBuf"); + numBytesAndAlignmentPerSmemChunk.emplace_back( + std::make_pair(numSmemBytesConstSfBuf, numBytesAlignmentConstSfBuf)); + firstChunkReuseSmem.emplace_back(reuseChunksSmemConstSfBuf); + } + + // Create SMEM helper object. + mSmemAllocatorHelper = + MemAllocatorHelper(numBytesAndAlignmentPerSmemChunk, firstChunkReuseSmem, smemChunkNames); +#if 0 + // E.g., + // Chunk 0 smemLoadA: 32768 bytes, 1024 alignment, false, offset 0 + // Chunk 1 smemLoadB: 32768 bytes, 1024 alignment, false, offset 32768 + // Chunk 2 smemBShuffle: 0 bytes, 1024 alignment, false, offset 65536 + // Chunk 3 smemGmemC0: 65536 bytes, 1024 alignment, true, offset 0 + // Chunk 4 smemGmemC1: 65536 bytes, 1024 alignment, false, offset 65536 + // Chunk 5 smemRowMax: 512 bytes, 16 alignment, false, offset 131072 + // Chunk 6 smemSliceK: 0 bytes, 16 alignment, false, offset 131584 + // Chunk 7 smemPerTokenSf: 0 bytes, 16 alignment, false, offset 131584 + mSmemAllocatorHelper.print(); +#endif + } + + // + // TMEM + // + // [..D..][..A..][.SfA.][.SfB.] + { + std::vector> numBytesAndAlignmentPerTmemChunk; + std::vector firstChunkReuseTmem; + std::vector tmemChunkNames; + // Matrix D + { + // Number of columns for accumulators. + auto const numTmemColsD = numSlicesForSliceK * tileN * numStagesMma * + tg::dtypeGetNumBits(dtypeAcc) / + tg::dtypeGetNumBits(tg::Dtype::UInt32); + // Number of columns for D alignment. + auto const numColsAlignmentD = 2; + // No need to reuse TMEM. + auto const reuseChunksTmemD = false; + + // Add info. + tmemChunkNames.emplace_back("tmemD"); + numBytesAndAlignmentPerTmemChunk.emplace_back( + std::make_pair(numTmemColsD, numColsAlignmentD)); + firstChunkReuseTmem.emplace_back(reuseChunksTmemD); + } + + // Matrix A + { + // We use TMEM for A if we use slice-K or if we need to cast A. + bool const useTmemA = (numSlicesForSliceK > 1) || (dtypeMmaA != dtypeA); + // Number of columns for A. + auto const numTmemColsA = + useTmemA ? numStages * tileK / + (numSlicesForSliceK * tg::dtypeGetNumBits(tg::Dtype::UInt32) / + tg::dtypeGetNumBits(dtypeMmaA)) + : 0; + // Number of columns for A alignment. + auto const numColsAlignmentA = 4; + // No need to reuse TMEM. + auto const reuseChunksTmemA = false; + + // Add info. + tmemChunkNames.emplace_back("tmemA"); + numBytesAndAlignmentPerTmemChunk.emplace_back( + std::make_pair(numTmemColsA, numColsAlignmentA)); + firstChunkReuseTmem.emplace_back(reuseChunksTmemA); + } + + // Sf A + { + // Does the MMA require block scales in TMEM for A? + bool const useBlockScalingA = tg::dtypeIsBlockFmt(dtypeMmaA); + // Are the block scales constant? + bool const useConstSfA = useBlockScalingA && !tg::dtypeIsBlockFmt(dtypeA); + // Number of columns for scaling factors of A. + auto const numTmemColsSfA = + useConstSfA + ? tg::roundUp((tileK / 64) * 2 * tg::ceilDiv(tileM, 64), 4) + : (useBlockScalingA ? ((tileK / 64) * 2 * tg::ceilDiv(tileM, 64)) * numStages : 0); + // Number of columns for Sf alignment. + auto const numColsAlignmentSfA = 4; + // No need to reuse TMEM. + auto const reuseChunksTmemSfA = false; + + // Add info. + tmemChunkNames.emplace_back("tmemSfA"); + numBytesAndAlignmentPerTmemChunk.emplace_back( + std::make_pair(numTmemColsSfA, numColsAlignmentSfA)); + firstChunkReuseTmem.emplace_back(reuseChunksTmemSfA); + } + + // Sf B + { + // Does the MMA require block scales in TMEM for B? + bool const useBlockScalingB = tg::dtypeIsBlockFmt(dtypeMmaB); + // Are the block scales constant? + bool const useConstSfB = useBlockScalingB && !tg::dtypeIsBlockFmt(dtypeB); + // Number of columns for scaling factors of B. + auto const numTmemColsSfB = + useConstSfB + ? tg::roundUp((tileK / 64) * 2 * tg::ceilDiv(tileN, 64), 4) + : (useBlockScalingB ? ((tileK / 64) * 2 * tg::ceilDiv(tileN, 64)) * numStages : 0); + // Number of columns for Sf alignment. + auto const numColsAlignmentSfB = 4; + // No need to reuse TMEM. + auto const reuseChunksTmemSfB = false; + + // Add info. + tmemChunkNames.emplace_back("tmemSfB"); + numBytesAndAlignmentPerTmemChunk.emplace_back( + std::make_pair(numTmemColsSfB, numColsAlignmentSfB)); + firstChunkReuseTmem.emplace_back(reuseChunksTmemSfB); + } + + // Create TMEM helper object. + mTmemAllocatorHelper = + MemAllocatorHelper(numBytesAndAlignmentPerTmemChunk, firstChunkReuseTmem, tmemChunkNames); + } + } + + public: + // The MMA kind. + tg::MmaKind mMmaKind; + // Helper for SMEM allocation. + MemAllocatorHelper mSmemAllocatorHelper; + // Helper for TMEM allocation. + MemAllocatorHelper mTmemAllocatorHelper; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline int32_t getSmemBufferSize(KernelTraits traits) { + return traits.mSmemAllocatorHelper.getTotalSize(); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline int32_t getTmemBufferSize(KernelTraits traits) { + return traits.mTmemAllocatorHelper.getTotalSize(); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// +// +// Starting address of each SMEM buffer. +// +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline int32_t getSmemOffsetLoadA(KernelTraits traits) { + return traits.mSmemAllocatorHelper.getChunkOffsetByName("smemLoadA"); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline int32_t getSmemOffsetLoadB(KernelTraits traits) { + return traits.mSmemAllocatorHelper.getChunkOffsetByName("smemLoadB"); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline int32_t getSmemOffsetLoadAb(KernelTraits traits) { return getSmemOffsetLoadA(traits); } + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline int32_t getSmemOffsetLoadShuffleB(KernelTraits traits) { + return traits.mSmemAllocatorHelper.getChunkOffsetByName("smemBShuffle"); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline int32_t getSmemOffsetGmemC(KernelTraits traits, int resIdx = 0) { + return traits.mSmemAllocatorHelper.getChunkOffsetByName("smemGmemC" + std::to_string(resIdx)); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline int32_t getSmemOffsetRowMax(KernelTraits traits) { + return traits.mSmemAllocatorHelper.getChunkOffsetByName("smemRowMax"); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline int32_t getSmemOffsetSliceK(KernelTraits traits) { + return traits.mSmemAllocatorHelper.getChunkOffsetByName("smemSliceK"); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline int32_t getSmemOffsetPerTokenSf(KernelTraits traits) { + return traits.mSmemAllocatorHelper.getChunkOffsetByName("smemPerTokenSf"); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline int32_t getSmemOffsetBias(KernelTraits traits) { + return traits.mSmemAllocatorHelper.getChunkOffsetByName("smemBias"); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline int32_t getSmemOffsetBlockAmax(KernelTraits traits) { + return traits.mSmemAllocatorHelper.getChunkOffsetByName("smemBlockAmax"); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline int32_t getSmemOffsetConstSfBuf(KernelTraits traits) { + return traits.mSmemAllocatorHelper.getChunkOffsetByName("smemConstSfBuf"); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline int32_t isSmemAbRepurposedToGmemC(KernelTraits traits, int resIdx = 0) { + return traits.mSmemAllocatorHelper.getFirstChunkReuseFlagByName("smemGmemC" + + std::to_string(resIdx)); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// +// +// Starting address of each TMEM buffer. +// +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline int32_t getTmemOffsetD(KernelTraits traits) { + return traits.mTmemAllocatorHelper.getChunkOffsetByName("tmemD"); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline int32_t getTmemOffsetA(KernelTraits traits) { + return traits.mTmemAllocatorHelper.getChunkOffsetByName("tmemA"); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline int32_t getTmemOffsetSfA(KernelTraits traits) { + return traits.mTmemAllocatorHelper.getChunkOffsetByName("tmemSfA"); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline int32_t getTmemOffsetSfB(KernelTraits traits) { + return traits.mTmemAllocatorHelper.getChunkOffsetByName("tmemSfB"); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace gemm + +} // namespace batchedGemm diff --git a/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/TmaDescriptor.h b/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/TmaDescriptor.h new file mode 100644 index 0000000000000000000000000000000000000000..a1412444aead854ac2a7f5f7f8f52194042d6415 --- /dev/null +++ b/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/TmaDescriptor.h @@ -0,0 +1,297 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & + * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include + +#include "trtllm/gen/DtypeDecl.h" +#include "trtllm/gen/MmaDecl.h" + +#ifdef TLLM_ENABLE_CUDA +#include +#include +#include +#endif + +namespace batchedGemm { + +namespace gemm { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace tg = trtllm::gen; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +#ifdef TLLM_ENABLE_CUDA + +inline CUtensorMap buildNdTmaDescriptor(tg::Dtype dtype, tg::MmaKind mmaKind, + std::vector const& shapes, + std::vector const& strides, + std::vector const& tileShapes, void* gmemAddr, + bool doSwizzle = true) { + // The multiplication factor of the data padding in SMEM. + int32_t padMultiplier = 1; + CUtensorMap desc{}; + // The data type. + CUtensorMapDataType tmaDataFormat{CU_TENSOR_MAP_DATA_TYPE_FLOAT32}; + if (dtype == tg::Dtype::E4m3 || dtype == tg::Dtype::MxE4m3 || dtype == tg::Dtype::UE8m0) { + tmaDataFormat = CU_TENSOR_MAP_DATA_TYPE_UINT8; + } else if (dtype == tg::Dtype::Fp16) { + tmaDataFormat = CU_TENSOR_MAP_DATA_TYPE_FLOAT16; + } else if (dtype == tg::Dtype::Bfloat16) { + tmaDataFormat = CU_TENSOR_MAP_DATA_TYPE_BFLOAT16; + } else if (dtype == tg::Dtype::E2m1) { + tmaDataFormat = CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN8B; + } else if (dtype == tg::Dtype::MxE2m1) { + if (mmaKind == tg::MmaKind::MxFp8Fp6Fp4) { + padMultiplier = 2; + tmaDataFormat = CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN16B; + } else { + // Note: this is used with the MMA kind MxFp4NvFp4 and also when casting to a higher-precision + // type such as Bfloat16 before the MMA. + tmaDataFormat = CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN8B; + } + } else if (dtype == tg::Dtype::Fp32) { + tmaDataFormat = CU_TENSOR_MAP_DATA_TYPE_FLOAT32; + } else { + std::cerr << "buildNdTmaDescriptor: unexpected dtype " << tg::dtypeToString(dtype) << std::endl; + assert(false); + } + + // The swizzle type. + CUtensorMapSwizzle swizzleType{CU_TENSOR_MAP_SWIZZLE_NONE}; + int32_t fastestDimTileSizeBytes = + (tileShapes[0] * tg::dtypeGetNumBits(dtype) * padMultiplier) / /* bits */ 8; + if (doSwizzle) { + if ((fastestDimTileSizeBytes % 128) == 0) { + swizzleType = CU_TENSOR_MAP_SWIZZLE_128B; + } else if ((fastestDimTileSizeBytes % 64) == 0) { + swizzleType = CU_TENSOR_MAP_SWIZZLE_64B; + } else if ((fastestDimTileSizeBytes % 32) == 0) { + swizzleType = CU_TENSOR_MAP_SWIZZLE_32B; + // This path is only for the scaling factors. + } else if ((fastestDimTileSizeBytes % 16) == 0 && + (dtype == tg::Dtype::UE8m0 || dtype == tg::Dtype::E4m3)) { + swizzleType = CU_TENSOR_MAP_SWIZZLE_NONE; + } else { + std::cerr << "buildNdTmaDescriptor: unexpected fastestDimTileSizeBytes " + << fastestDimTileSizeBytes << std::endl; + assert(false); + } + } + + // Check gmem address must be 16B-aligned + assert((reinterpret_cast(gmemAddr) & 0b1111) == 0); // + + // Check shape must be in range [1, 2^32] + int32_t dim = shapes.size(); + // Expect 2 dimensions for regular gemm, 3 dimensions for batched gemm or blocked layout, and 4 + // dimensions for batched gemm with blocked layout. + assert(dim == 2 || dim == 3 || dim == 4); + // Check shape range. + for (int32_t ii = 0; ii < dim; ++ii) { + assert(shapes[ii] >= (uint64_t(1))); // Size must be min 1 + assert(shapes[ii] <= (uint64_t(1) << 32)); // Size must be max 2^32 + } + + // TMA descriptor does not store the zeroth stride and assumes it is 1. + assert(static_cast(strides.size()) == dim); + assert(strides[0] == 1); + + // Build strides in bytes. + // cuTensorMapEncodeTiled ignores the stride of the first dimension (implicitly 1). + std::vector stridesInBytes(dim - 1); + for (int32_t ii = 0; ii < dim - 1; ++ii) { + stridesInBytes[ii] = (strides[ii + 1] * tg::dtypeGetNumBits(dtype)) / /* bits */ 8; + } + + // Set the number of elements in the packed uint32_t element. + auto const numEltsPerUInt32 = 4 * /* bits */ 8 / (tg::dtypeGetNumBits(dtype) * padMultiplier); + // The number of elements in 128B. + auto const numEltsIn128B = numEltsPerUInt32 /*4B*/ * 32; + // The number of tile K hidden size (per token) in each block of shared memory. + auto const numEltsInClampedFastestTileSize = std::min(numEltsIn128B, tileShapes[0]); + + // Build box dim array. If tileShapes is smaller than dim, just fill with 1s. + assert(static_cast(tileShapes.size()) <= dim); + std::vector boxDim(dim, 1); + boxDim[0] = numEltsInClampedFastestTileSize; + for (size_t ii = 1; ii < tileShapes.size(); ++ii) { + if (tileShapes[ii] > 256) { + std::cerr << "buildNdTmaDescriptor: boxDim too large " << tileShapes[ii] << std::endl; + assert(false); + } else { + boxDim[ii] = tileShapes[ii]; + } + } + + // Set tile strides to 1; + std::vector tileStrides(dim, 1); + + // Build the descriptor. + CUresult result = + cuTensorMapEncodeTiled(&desc, tmaDataFormat, + /*tensorRank=*/dim, gmemAddr, shapes.data(), stridesInBytes.data(), + boxDim.data(), tileStrides.data(), + /*interleave=*/CU_TENSOR_MAP_INTERLEAVE_NONE, swizzleType, + /*l2Promotion=*/CU_TENSOR_MAP_L2_PROMOTION_L2_128B, + /*oobFill=*/CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + + if (result != CUDA_SUCCESS) { + char const* errorString; + cuGetErrorString(result, &errorString); + std::stringstream ss; + ss << "Error: Failed to initialize the TMA descriptor " << result << std::endl; + + ss << "tmaFormat: " << static_cast(tmaDataFormat) << " dim: " << dim + << " gmem: " << gmemAddr << std::endl; + + ss << "Shape: "; + for (int ii = 0; ii < dim; ++ii) { + ss << shapes[ii] << " "; + } + ss << std::endl; + + ss << "Stride: "; + for (int ii = 0; ii < dim - 1; ++ii) { + ss << stridesInBytes[ii] << " "; + } + ss << std::endl; + + ss << "tileShapes: "; + for (int ii = 0; ii < dim; ++ii) { + ss << boxDim[ii] << " "; + } + ss << std::endl; + + ss << "tileStrides: "; + for (int ii = 0; ii < dim; ++ii) { + ss << tileStrides[ii] << " "; + } + ss << std::endl; + ss << "swizzleType: " << int(swizzleType) << std::endl; + ss << "(in " << __FILE__ << ":" << __LINE__ << ")" << std::endl; + throw std::runtime_error(ss.str()); + } + + return desc; +} + +// TODO: make it work with the above descriptor? +inline CUtensorMap buildSfTmaDescriptor(tg::Dtype dtype, std::vector const& shapes, + std::vector const& strides, + const std::vector& tileShapes, void* gmemAddr) { + CUtensorMap desc{}; + CUtensorMapDataType tmaDataFormat; + if (dtype == tg::Dtype::E4m3 || dtype == tg::Dtype::UE8m0) { + tmaDataFormat = CU_TENSOR_MAP_DATA_TYPE_UINT8; + } else { + std::cerr << "buildSfTmaDescriptor: unexpected dtype " << tg::dtypeToString(dtype) << std::endl; + assert(false); + } + + // No swizzle for scaling factors. + CUtensorMapSwizzle swizzleType = CU_TENSOR_MAP_SWIZZLE_NONE; + + // Check gmem address must be 16B-aligned + assert((reinterpret_cast(gmemAddr) & 0b1111) == 0); // + + // Check shape must be in range [1, 2^32] + int32_t dim = shapes.size(); + // Check shape range. + for (int32_t ii = 0; ii < dim; ++ii) { + assert(shapes[ii] >= (uint64_t(1))); // Size must be min 1 + assert(shapes[ii] <= (uint64_t(1) << 32)); // Size must be max 2^32 + } + + // TMA descriptor does not store the zeroth stride and assumes it is 1. + assert(static_cast(strides.size()) == dim); + assert(strides[0] == 1); + + // Build strides in bytes. + // cuTensorMapEncodeTiled ignores the stride of the first dimension (implicitly 1). + std::vector stridesInBytes(dim - 1); + for (int32_t ii = 0; ii < dim - 1; ++ii) { + stridesInBytes[ii] = (strides[ii + 1] * tg::dtypeGetNumBits(dtype)) / /* bits */ 8; + } + + // Set tile strides to 1; + std::vector tileStrides(dim, 1); + + // Build the descriptor. + CUresult result = cuTensorMapEncodeTiled(/*tensorMap=*/&desc, + /*tensorDataType=*/tmaDataFormat, + /*tensorRank=*/dim, + /*globalAddress=*/gmemAddr, + /*globalDim=*/shapes.data(), + /*globalStrides=*/stridesInBytes.data(), + /*boxDim=*/tileShapes.data(), + /*elementStrides=*/tileStrides.data(), + /*interleave=*/CU_TENSOR_MAP_INTERLEAVE_NONE, + /*swizzle=*/swizzleType, + /*l2Promotion=*/CU_TENSOR_MAP_L2_PROMOTION_L2_128B, + /*oobFill=*/CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + + if (result != CUDA_SUCCESS) { + char const* errorString; + cuGetErrorString(result, &errorString); + std::stringstream ss; + ss << "Error: Failed to initialize the TMA descriptor for SF " << errorString << std::endl; + + ss << "tmaFormat: " << static_cast(tmaDataFormat) << " dim: " << dim + << " gmem: " << gmemAddr << std::endl; + + ss << "shape:"; + for (uint32_t shape_i : shapes) { + ss << " " << shape_i; + } + ss << std::endl; + + ss << "stridesInBytes:"; + for (uint32_t stride_i : stridesInBytes) { + ss << " " << stride_i; + } + ss << std::endl; + + ss << "tileShapes:"; + for (uint32_t tileShape_i : tileShapes) { + ss << " " << tileShape_i; + } + ss << std::endl; + + ss << "tileStrides:"; + for (uint32_t tileStride_i : tileStrides) { + ss << " " << tileStride_i; + } + ss << std::endl; + + ss << "swizzleType: " << int(swizzleType) << std::endl; + ss << "(in " << __FILE__ << ":" << __LINE__ << ")" << std::endl; + throw std::runtime_error(ss.str()); + } + + return desc; +} + +#endif // defined TLLM_ENABLE_CUDA + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace gemm + +} // namespace batchedGemm diff --git a/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/CommonUtils.h b/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/CommonUtils.h new file mode 100644 index 0000000000000000000000000000000000000000..393949a5165243ae951f17f4506bbb42763e669a --- /dev/null +++ b/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/CommonUtils.h @@ -0,0 +1,43 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & + * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +namespace batchedGemm { + +namespace trtllm { +namespace gen { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +inline T ceilDiv(T m, T n) { + return (m + n - T(1)) / n; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +inline T roundUp(T m, T n) { + return ceilDiv(m, n) * n; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace gen +} // namespace trtllm + +} // namespace batchedGemm diff --git a/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/CudaKernelLauncher.h b/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/CudaKernelLauncher.h new file mode 100644 index 0000000000000000000000000000000000000000..42bc884f925042babdfa85b363099bd3518b8ef5 --- /dev/null +++ b/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/CudaKernelLauncher.h @@ -0,0 +1,92 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & + * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#ifdef TLLM_ENABLE_CUDA +#include +#include + +#include +#include +#endif +namespace batchedGemm { + +namespace trtllm { +namespace gen { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +#ifdef TLLM_ENABLE_CUDA +inline CUresult launchKernel(void* kernelParams, void* cudaStream, int32_t smemSize, + CUfunction kernel, dim3 block3, dim3 grid3, dim3 cluster3, + bool enablesPdl) { + // Make sure we can launch with that much shared memory. + if (smemSize > 48 * 1024) { + CUresult result = + cuFuncSetAttribute(kernel, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, smemSize); + if (result != CUDA_SUCCESS) { + return result; + } + } + + auto clusterDim = cluster3.x * cluster3.y * cluster3.z; + + CUlaunchConfig launchConfig; + launchConfig.blockDimX = block3.x; + launchConfig.blockDimY = block3.y; + launchConfig.blockDimZ = block3.z; + launchConfig.gridDimX = grid3.x; + launchConfig.gridDimY = grid3.y; + launchConfig.gridDimZ = grid3.z; + launchConfig.hStream = reinterpret_cast(cudaStream); + launchConfig.sharedMemBytes = smemSize; + + CUlaunchAttribute launchAttrs[3]; + launchAttrs[0].id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION; + launchAttrs[0].value.clusterDim.x = cluster3.x; + launchAttrs[0].value.clusterDim.y = cluster3.y; + launchAttrs[0].value.clusterDim.z = cluster3.z; + launchAttrs[1].id = CU_LAUNCH_ATTRIBUTE_CLUSTER_SCHEDULING_POLICY_PREFERENCE; + launchAttrs[1].value.clusterSchedulingPolicyPreference = + (clusterDim > 1) ? CU_CLUSTER_SCHEDULING_POLICY_SPREAD : CU_CLUSTER_SCHEDULING_POLICY_DEFAULT; + launchAttrs[2].id = CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION; + launchAttrs[2].value.programmaticStreamSerializationAllowed = enablesPdl; + launchConfig.attrs = launchAttrs; + launchConfig.numAttrs = 3; + + // Add setting for non-portable cluster size. + if (clusterDim > 8) { + CUresult result = + cuFuncSetAttribute(kernel, CU_FUNC_ATTRIBUTE_NON_PORTABLE_CLUSTER_SIZE_ALLOWED, + 1 // Enable non-portable cluster sizes + ); + if (result != CUDA_SUCCESS) { + return result; + } + } + + // Launch the kernel. + return cuLaunchKernelEx(&launchConfig, kernel, &kernelParams, nullptr); +} +#endif + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace gen +} // namespace trtllm + +} // namespace batchedGemm diff --git a/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/DtypeDecl.h b/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/DtypeDecl.h new file mode 100644 index 0000000000000000000000000000000000000000..0866256492c74a03d65a4c731ee9cdeebd1c29e7 --- /dev/null +++ b/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/DtypeDecl.h @@ -0,0 +1,271 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & + * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include +#include +#include +#include +#ifndef TLLM_GEN_EXPORT_INTERFACE +#include "trtllm/gen/MmaDecl.h" +#else +#include "MmaDecl.h" +#endif + +//////////////////////////////////////////////////////////////////////////////////////////////////// +// +// Be careful when modifying this file as it is included by the generated kernels. For example, do +// not add TLLM_CHECK_* constructs in this file. Thanks! +// +//////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace batchedGemm { + +namespace trtllm { +namespace gen { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +enum class Dtype : uint32_t { + +// We use the following encoding for the types: +// +// Byte 0: Identifier for the type (going from 0 to the number of data types - 1, +// Byte 1: Number of bits in the type, +// Byte 2: Bit 0: Is it an integer? 0x1 if true, 0x0 otherwise; +// Bit 4: is it signed? 0x1 if true, 0x0 otherwise. +// Byte 3: Is it a block format? 0x1 if true, 0x0 otherwise. + +#define TLLM_ENCODE_DTYPE(BlockFormatBit, SignedBit, IntegerBit, NumBits, Uid) \ + uint32_t { \ + (BlockFormatBit << 24) | (SignedBit << 20) | (IntegerBit << 16) | (NumBits << 8) | (Uid) \ + } + + // clang-format off + Bfloat16 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 1u, /*int*/ 0u, /*bits*/ 16u, /*uid*/ 0u), + Bool = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 0u, /*int*/ 1u, /*bits*/ 1u, /*uid*/ 1u), + E2m1 = TLLM_ENCODE_DTYPE(/*block*/ 1u, /*signed*/ 1u, /*int*/ 0u, /*bits*/ 4u, /*uid*/ 2u), + E2m3 = TLLM_ENCODE_DTYPE(/*block*/ 1u, /*signed*/ 1u, /*int*/ 0u, /*bits*/ 6u, /*uid*/ 3u), + E3m2 = TLLM_ENCODE_DTYPE(/*block*/ 1u, /*signed*/ 1u, /*int*/ 0u, /*bits*/ 6u, /*uid*/ 4u), + E4m3 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 1u, /*int*/ 0u, /*bits*/ 8u, /*uid*/ 5u), + E5m2 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 1u, /*int*/ 0u, /*bits*/ 8u, /*uid*/ 6u), + Fp16 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 1u, /*int*/ 0u, /*bits*/ 16u, /*uid*/ 7u), + Fp32 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 1u, /*int*/ 0u, /*bits*/ 32u, /*uid*/ 8u), + Int8 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 1u, /*int*/ 1u, /*bits*/ 8u, /*uid*/ 9u), + Int32 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 1u, /*int*/ 1u, /*bits*/ 32u, /*uid*/ 10u), + Int64 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 1u, /*int*/ 1u, /*bits*/ 64u, /*uid*/ 11u), + MxE2m1 = TLLM_ENCODE_DTYPE(/*block*/ 1u, /*signed*/ 1u, /*int*/ 0u, /*bits*/ 4u, /*uid*/ 12u), + MxE4m3 = TLLM_ENCODE_DTYPE(/*block*/ 1u, /*signed*/ 1u, /*int*/ 0u, /*bits*/ 8u, /*uid*/ 13u), + UE8m0 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 0u, /*int*/ 0u, /*bits*/ 8u, /*uid*/ 14u), + UInt8 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 0u, /*int*/ 1u, /*bits*/ 8u, /*uid*/ 15u), + UInt16 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 0u, /*int*/ 1u, /*bits*/ 16u, /*uid*/ 16u), + UInt32 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 0u, /*int*/ 1u, /*bits*/ 32u, /*uid*/ 17u), + UInt64 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 0u, /*int*/ 1u, /*bits*/ 64u, /*uid*/ 18u), + UInt128 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 0u, /*int*/ 1u, /*bits*/ 128u, /*uid*/ 19u), + Void = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 1u, /*int*/ 0u, /*bits*/ 0u, /*uid*/ 20u), +// clang-format on + +#undef TLLM_ENCODE_DTYPE +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// The number of bits in a data type? +inline int dtypeGetNumBits(Dtype dtype) { + constexpr uint32_t kMask = 0xffu << 8; + return static_cast((static_cast(dtype) & kMask) >> 8); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// Does the format use block scaling? +inline bool dtypeIsBlockFmt(Dtype dtype) { + constexpr uint32_t kMask = 0xffu << 24; + return static_cast((static_cast(dtype) & kMask) >> 24); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// Is a given data type a floating-point type? +inline bool dtypeIsFloat(Dtype dtype) { + constexpr uint32_t kMask = 0x1u << 16; + return dtype != Dtype::Void && 0 == (static_cast(dtype) & kMask); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// Is a given data type an 8-bit floating-point type? +inline bool dtypeIsFp8(Dtype dtype) { return dtype == Dtype::E4m3 || dtype == Dtype::E5m2; } + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// Is a given data type an integer type? +inline bool dtypeIsInt(Dtype dtype) { + constexpr uint32_t kMask = 0x1u << 16; + return (dtype != Dtype::Bool) && (0 != (static_cast(dtype) & kMask)); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// Is a given data type signed? +inline bool dtypeIsSigned(Dtype dtype) { + constexpr uint32_t kMask = 0x1u << 20; + return (0 != (static_cast(dtype) & kMask)); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// For logging and error reporting +inline std::string dtypeToString(Dtype dtype) { + switch (dtype) { + case Dtype::Bfloat16: + return "Bfloat16"; + case Dtype::Bool: + return "Bool"; + case Dtype::E2m1: + return "E2m1"; + case Dtype::E2m3: + return "E2m3"; + case Dtype::E3m2: + return "E3m2"; + case Dtype::E4m3: + return "E4m3"; + case Dtype::E5m2: + return "E5m2"; + case Dtype::Fp16: + return "Fp16"; + case Dtype::Fp32: + return "Fp32"; + case Dtype::Int8: + return "Int8"; + case Dtype::Int32: + return "Int32"; + case Dtype::Int64: + return "Int64"; + case Dtype::MxE4m3: + return "MxE4m3"; + case Dtype::MxE2m1: + return "MxE2m1"; + case Dtype::UE8m0: + return "UE8m0"; + case Dtype::UInt8: + return "UInt8"; + case Dtype::UInt16: + return "UInt16"; + case Dtype::UInt32: + return "UInt32"; + case Dtype::UInt64: + return "UInt64"; + case Dtype::UInt128: + return "UInt128"; + case Dtype::Void: + return "Void"; + default: + assert(false); + return "Unsupported type"; + } +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline Dtype dtypeEltType(Dtype dtype) { + switch (dtype) { + case Dtype::MxE2m1: + return Dtype::E2m1; + case Dtype::MxE4m3: + return Dtype::E4m3; + default: + return dtype; + } +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline int dtypeNumEltsPerSf(Dtype dtype) { + switch (dtype) { + case Dtype::E2m1: + return 16; + case Dtype::MxE2m1: + case Dtype::MxE4m3: + return 32; + default: + assert(false); + return -1; + } +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// Returns the dtype of scaling factors, if applicable. +inline Dtype dtypeGetBlockSfType(Dtype dtype) { + switch (dtype) { + case Dtype::E2m1: + return Dtype::E4m3; + case Dtype::MxE2m1: + case Dtype::MxE4m3: + return Dtype::UE8m0; + default: + assert(false); + return Dtype::Void; + } +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline MmaKind dtypeGetMmaKind(Dtype dtypeA, Dtype dtypeB) { + auto dtypeEltA = dtypeEltType(dtypeA); + auto dtypeEltB = dtypeEltType(dtypeB); + + // Note: the order of the conditions is important here. + if ((dtypeA == Dtype::Fp16 && dtypeB == Dtype::Fp16) || + (dtypeA == Dtype::Bfloat16 && dtypeB == Dtype::Bfloat16)) { + return MmaKind::Fp16; + } + + if ((dtypeA == Dtype::Int8 || dtypeA == Dtype::UInt8) && + (dtypeB == Dtype::Int8 || dtypeB == Dtype::UInt8)) { + return MmaKind::Int8; + } + + // This statement captures both MxE2m1 and E2m1. + if (dtypeEltA == Dtype::E2m1 && dtypeEltB == Dtype::E2m1) { + return MmaKind::MxFp4NvFp4; + } + + if ((dtypeA == Dtype::E4m3 || dtypeA == Dtype::E5m2 || dtypeA == Dtype::E2m3 || + dtypeA == Dtype::E3m2 || dtypeA == Dtype::E2m1) && + (dtypeB == Dtype::E4m3 || dtypeB == Dtype::E5m2 || dtypeB == Dtype::E2m3 || + dtypeB == Dtype::E3m2 || dtypeB == Dtype::E2m1)) { + return MmaKind::Fp8Fp6Fp4; + } + + // At this point we know that both dtypes are Mx types and not both MxE2m1 at the same time. + if ((dtypeEltA == Dtype::E4m3 || dtypeEltA == Dtype::E5m2 || dtypeEltA == Dtype::E2m3 || + dtypeEltA == Dtype::E3m2 || dtypeEltA == Dtype::E2m1) && + (dtypeEltB == Dtype::E4m3 || dtypeEltB == Dtype::E5m2 || dtypeEltB == Dtype::E2m3 || + dtypeEltB == Dtype::E3m2 || dtypeEltB == Dtype::E2m1)) { + return MmaKind::MxFp8Fp6Fp4; + } + return MmaKind::Tf32; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace gen +} // namespace trtllm + +} // namespace batchedGemm diff --git a/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/MmaDecl.h b/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/MmaDecl.h new file mode 100644 index 0000000000000000000000000000000000000000..27f5e36d9c96d9cf0e5dcdbe868d27c3b3ef2a2b --- /dev/null +++ b/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/MmaDecl.h @@ -0,0 +1,92 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & + * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +namespace batchedGemm { + +namespace trtllm { +namespace gen { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// The kind of the MMA instruction +enum class MmaKind : uint32_t { + // For Blackwell this follows the PTX ISA description of the MMA instructions. + // https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-kind-shapes + + // The MMA type is auto-detected from the dtypes of the input tensors + Auto = 0, + // Supports dtypeA = dtypeB = Fp16 and dtypeD = [Fp16, Fp32] + // or dtypeA = dtypeB = Bfloat16 and dtypeD = [Fp32] + // Corresponds to the kind::f16 of tcgen05.mma. + Fp16 = 1, + // Supports dtypeA/B = [E4m3, E5m2, E2m3, E3m2, E2m1] and dtypeD = [Fp16, Fp32] + // Corresponds to the kind::f8f6f4 of tcgen05.mma. + Fp8Fp6Fp4 = 2, + // Supports dtypeA = dtypeB = [Int8, Uint8] and dtypeD = [Int32] + // Corresponds to the kind::i8 of tcgen05.mma. + Int8 = 3, + // Supports dtypeA = dtypeB = [MxE2m1, E2m1] with block scale [UM8e0, UEm4e3] + // and dtypeD = [Fp32] + // Corresponds to the kind::mxf4nvf4 of tcgen05.mma. + MxFp4NvFp4 = 4, + // Supports dtype dtypeA = dtypeB = [MxE4m3, MxE2m1] with block scale [UM8e0] + // and dtypeD = [Fp32] + // Corresponds to the kind::mxf8f6f4 of tcgen05.mma. + MxFp8Fp6Fp4 = 5, + // Supports dtypeA = dtypeB = Tf32 with dtypeD = [Fp32] + // Corresponds to the kind::tf32 of tcgen05.mma. + Tf32 = 6 +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline bool mmaKindIsBlockFmt(MmaKind mmaKind) { + return mmaKind == MmaKind::MxFp8Fp6Fp4 || mmaKind == MmaKind::MxFp4NvFp4; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// For logging and error reporting +inline std::string mmaKindToString(MmaKind mmaKind) { + switch (mmaKind) { + case MmaKind::Auto: + return "Auto"; + case MmaKind::Fp16: + return "Fp16"; + case MmaKind::Fp8Fp6Fp4: + return "Fp8Fp6Fp4"; + case MmaKind::Int8: + return "Int8"; + case MmaKind::MxFp4NvFp4: + return "MxFp4NvFp4"; + case MmaKind::MxFp8Fp6Fp4: + return "MxFp8Fp6Fp4"; + case MmaKind::Tf32: + return "Tf32"; + default: + assert(false); + return "Unsupported type"; + } +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace gen +} // namespace trtllm + +} // namespace batchedGemm diff --git a/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/SfLayoutDecl.h b/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/SfLayoutDecl.h new file mode 100644 index 0000000000000000000000000000000000000000..965bb1b7b84511994aa99f0587bf6392b7ed11bf --- /dev/null +++ b/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/SfLayoutDecl.h @@ -0,0 +1,96 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & + * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include +#include + +//////////////////////////////////////////////////////////////////////////////////////////////////// +// +// Be careful when modifying this file as it is included by the generated kernels. For example, do +// not add TLLM_CHECK_* constructs in this file. Thanks! +// +//////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace batchedGemm { + +namespace trtllm { +namespace gen { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// This enumeration defines layouts for storing scale factors for FP4, FP6, and FP8 formats. +enum class SfLayout { + // Scale factors are stored in the same order as the associated matrix. + // I.e., the SF buffer is a tensor [m, ⌈n/b⌉], where m, n, and b are respectively the number of + // rows, columns and the block size. + // The SF for the element (i, j) is stored at (i, j/b). + Linear = 0, + + // A tile of 8x4 is stored contiguously. The order of elements inside the tile, and the order + // of tiles, are both row-major. + // I.e., the SF buffer is a tensor [⌈m/8⌉, ⌈n/b/4⌉, 8, 4]. + // The SF for the element (i, j) is stored at (i/8, j/b/4, i%8, (j/b)%4). + R8c4, + + // A tile of 8x16 is stored contiguously. The order of elements inside the tile, and the order + // of tiles, are both row-major. + // I.e., the SF buffer is a tensor [⌈m/8⌉, ⌈n/b/16⌉, 8, 16]. + // The SF for the element (i, j) is stored at (i/8, j/b/16, i%8, (j/b)%16). + // + // NOTE: This is a niche format that is currently used for the weights of the + // LowLatency FP4 kernels. It is not meant as an interchange format. In + // addition to the above requirements it requires n to be a multiple of 256. + R8c16, + + // A tile of 128x4 is stored contiguously. Rows 0-31, 32-63, 64-95 and 96-127 are interleaved + // as illustrated below: + // | 0,0 | 0,1 | 0,2 | 0,3 | 32,0 | 32,1 | 32,2 | 32,3 | ... | 96,3 | + // | 1,0 | 1,1 | 1,2 | 1,3 | 33,0 | 33,1 | 33,2 | 33,3 | ... | 97,3 | + // | ... | ... | ... | ... | ... | ... | ... | ... | ... | ... | + // | 31,0 | 31,1 | 31,2 | 31,3 | 63,0 | 63,1 | 63,2 | 63,3 | ... | 127,3 | + // See https://nvbugspro.nvidia.com/bug/4165523 + // + // I.e., the SF buffer is a tensor [⌈m/128⌉, ⌈n/b/4⌉, 32, 4, 4] + // The SF for the element (i, j) is stored at (i/128, j/b/4, i%32, (i%128)/32, (j/b)%4). + R128c4, +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline std::string sfLayoutToString(SfLayout layout) { + switch (layout) { + case SfLayout::Linear: + return "linear"; + case SfLayout::R8c4: + return "8x4"; + case SfLayout::R8c16: + return "8x16"; + case SfLayout::R128c4: + return "128x4"; + default: + assert(false); + return "Unsupported layout"; + } +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace gen +} // namespace trtllm + +} // namespace batchedGemm diff --git a/include/flashinfer/trtllm/common.h b/include/flashinfer/trtllm/common.h new file mode 100644 index 0000000000000000000000000000000000000000..63abd839107a96a1774bbcd74c7850bf0c429b2d --- /dev/null +++ b/include/flashinfer/trtllm/common.h @@ -0,0 +1,335 @@ +/* + * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include +#include +#include +#include +#ifndef _WIN32 // Linux +#include +#endif // not WIN32 +#ifdef _WIN32 // Windows +#include +#undef ERROR // A Windows header file defines ERROR as 0, but it's used in our logger.h enum. + // Logging breaks without this undef. +#endif // WIN32 + +#include + +#define HOST_DEVICE_FUNC __host__ __device__ +#define DEVICE_FUNC __device__ + +inline void cuErrCheck_(CUresult stat, char const* file, int line) { + if (stat != CUDA_SUCCESS) { + char const* msg = nullptr; + cuGetErrorName(stat, &msg); + fprintf(stderr, "CUDA Error: %s %s %d\n", msg, file, line); + } +} +#define cuErrCheck(stat) \ + { \ + cuErrCheck_((stat), __FILE__, __LINE__); \ + } + +#define CUDACHECK(cmd) \ + do { \ + cudaError_t e = cmd; \ + if (e != cudaSuccess) { \ + printf("Failed: Cuda error %s:%d '%s'\n", __FILE__, __LINE__, cudaGetErrorString(e)); \ + exit(EXIT_FAILURE); \ + } \ + } while (0) + +inline constexpr int kMinHistoryTokensPerBlock = 128; + +inline constexpr float kEnableMinBlockFactor = 4.0; +inline constexpr int kTargetWaveFactor = 8; + +// For multi-block mode. We reserve workspace for this amount of sub-sequences. +// This should be enough. Huge batch size may result in larger value, but for +// large batch size, multi-block mode is not useful. For llama v2 70b, 6000 +// results in ~12MB multi-block workspace, and is enough for > 10 waves. +inline constexpr int kXQA_MAX_NUM_SUB_SEQ = 6000; +inline constexpr int kMaxBeamWidth = 1; + +inline int getDevice() { + int current_dev_id = 0; + CUDACHECK(cudaGetDevice(¤t_dev_id)); + return current_dev_id; +} +inline int getSMVersion() { + int device{-1}; + CUDACHECK(cudaGetDevice(&device)); + int sm_major = 0; + int sm_minor = 0; + CUDACHECK(cudaDeviceGetAttribute(&sm_major, cudaDevAttrComputeCapabilityMajor, device)); + CUDACHECK(cudaDeviceGetAttribute(&sm_minor, cudaDevAttrComputeCapabilityMinor, device)); + return sm_major * 10 + sm_minor; +} + +inline std::tuple getDeviceMemoryInfo(bool const useUvm) { + if (useUvm) { + size_t freeSysMem = 0; + size_t totalSysMem = 0; +#ifndef _WIN32 // Linux + struct sysinfo info{}; + + sysinfo(&info); + totalSysMem = info.totalram * info.mem_unit; + freeSysMem = info.freeram * info.mem_unit; +#else // Windows + MEMORYSTATUSEX memInfo; + memInfo.dwLength = sizeof(memInfo); + GlobalMemoryStatusEx(&memInfo); + totalSysMem = memInfo.ullTotalPhys; + freeSysMem = memInfo.ullAvailPhys; +#endif // WIN32 + + // printf("Using UVM based system memory for KV cache, total memory %0.2f GB, available memory + //%0.2f GB", + // ((double) totalSysMem / 1e9), ((double) freeSysMem / 1e9)); + // return {freeSysMem, totalSysMem}; + } + + size_t free = 0; + size_t total = 0; + CUDACHECK(cudaMemGetInfo(&free, &total)); + // printf("Using GPU memory for KV cache, total memory %0.2f GB, available memory %0.2f GB", + // ((double) total / 1e9), ((double) free / 1e9)); + return {free, total}; +} + +enum class LogLevel { + LOG_NULL, + LOG_ERROR, + LOG_WARNING, + LOG_INFO, + LOG_DEBUG, +}; + +inline LogLevel parseLogLevel(const char* env) { + if (env == nullptr) { + return LogLevel::LOG_NULL; // default + } + + if (std::strcmp(env, "ERROR") == 0) return LogLevel::LOG_ERROR; + if (std::strcmp(env, "WARNING") == 0) return LogLevel::LOG_WARNING; + if (std::strcmp(env, "INFO") == 0) return LogLevel::LOG_INFO; + if (std::strcmp(env, "DEBUG") == 0) return LogLevel::LOG_DEBUG; + + // Try numeric values as fallback + if (std::strcmp(env, "1") == 0) return LogLevel::LOG_ERROR; + if (std::strcmp(env, "2") == 0) return LogLevel::LOG_WARNING; + if (std::strcmp(env, "3") == 0) return LogLevel::LOG_INFO; + if (std::strcmp(env, "4") == 0) return LogLevel::LOG_DEBUG; + + return LogLevel::LOG_DEBUG; // default if unrecognized +} + +inline LogLevel getLogLevelFromEnv() { + static std::once_flag init_flag; + static LogLevel level = LogLevel::LOG_NULL; + + std::call_once(init_flag, []() { + const char* env = std::getenv("FLASHINFER_LOG_LEVEL"); + level = parseLogLevel(env); + }); + + return level; +} +// Global log level access +inline LogLevel getCurrentLogLevel() { return getLogLevelFromEnv(); } + +// Logging macros +#define IKL_LOG_DEBUG(fmt, ...) \ + do { \ + if (getCurrentLogLevel() >= LogLevel::LOG_DEBUG) { \ + printf("[DEBUG] " fmt "\n", ##__VA_ARGS__); \ + } \ + } while (0) + +#define IKL_LOG_INFO(fmt, ...) \ + do { \ + if (getCurrentLogLevel() >= LogLevel::LOG_INFO) { \ + printf("[INFO] " fmt "\n", ##__VA_ARGS__); \ + } \ + } while (0) + +#define IKL_LOG_WARNING(fmt, ...) \ + do { \ + if (getCurrentLogLevel() >= LogLevel::LOG_WARNING) { \ + printf("[WARNING] " fmt "\n", ##__VA_ARGS__); \ + } \ + } while (0) + +#define IKL_LOG_ERROR(fmt, ...) \ + do { \ + if (getCurrentLogLevel() >= LogLevel::LOG_ERROR) { \ + printf("[ERROR] " fmt "\n", ##__VA_ARGS__); \ + } \ + } while (0) + +// Returns true if the env variable exists and is set to "1" +inline static bool getBoolEnv(char const* name) { + char const* env = std::getenv(name); + return env && env[0] == '1' && env[1] == '\0'; +} + +inline bool getEnvUseTileSizeKv64ForTrtllmGen() { + static bool const useTileSizeKv64 = getBoolEnv("TRTLLM_GEN_ENABLE_TILE_SIZE_KV64"); + return useTileSizeKv64; +} + +inline bool getEnvEnablePDL() { + static std::once_flag flag; + static bool enablePDL = false; + + std::call_once(flag, [&]() { + if (getSMVersion() >= 90) { + // PDL will be enabled by setting the env variables `TRTLLM_ENABLE_PDL` to `1` + enablePDL = getBoolEnv("TRTLLM_ENABLE_PDL"); + } + }); + return enablePDL; +} +template +inline __device__ __host__ T divUp(T m, T n) { + return (m + n - 1) / n; +} + +// For gen kernel IO +enum Data_type { + DATA_TYPE_FP16, + DATA_TYPE_BF16, + DATA_TYPE_FP32, + DATA_TYPE_INT8, + DATA_TYPE_INT32, + DATA_TYPE_E4M3, + DATA_TYPE_E5M2, + DATA_TYPE_E2M1, + DATA_TYPE_UNKNOWN +}; + +inline constexpr const char* toStr(Data_type dtype) { + switch (dtype) { + case DATA_TYPE_FP16: + return "FP16"; + case DATA_TYPE_BF16: + return "BF16"; + case DATA_TYPE_FP32: + return "FP32"; + case DATA_TYPE_INT8: + return "INT8"; + case DATA_TYPE_INT32: + return "INT32"; + case DATA_TYPE_E4M3: + return "E4M3"; + case DATA_TYPE_E5M2: + return "E5M2"; + case DATA_TYPE_E2M1: + return "E2M1"; + default: + return "UNKNOWN"; + } +} + +// Type trait to map types to enum values +template +struct TypeToDataType { + static constexpr Data_type value = Data_type::DATA_TYPE_UNKNOWN; +}; + +// Specialize the trait for specific types +template <> +struct TypeToDataType<__nv_bfloat16> { + static constexpr Data_type value = Data_type::DATA_TYPE_BF16; +}; + +template <> +struct TypeToDataType<__half> { + static constexpr Data_type value = Data_type::DATA_TYPE_FP16; +}; + +template <> +struct TypeToDataType { + static constexpr Data_type value = Data_type::DATA_TYPE_E4M3; +}; + +template <> +struct TypeToDataType<__nv_fp8_e4m3> { + static constexpr Data_type value = Data_type::DATA_TYPE_E4M3; +}; + +static inline size_t get_size_in_bytes(size_t n, Data_type dtype) { + switch (dtype) { + case DATA_TYPE_FP32: + return n * 4; + case DATA_TYPE_FP16: + return n * 2; + case DATA_TYPE_INT32: + return n * 4; + case DATA_TYPE_INT8: + return n; + case DATA_TYPE_BF16: + return n * 2; + case DATA_TYPE_E4M3: + return n; + case DATA_TYPE_E5M2: + return n; + default: + TORCH_CHECK(false, "FMHA Data Type is not supported."); + return 0; + } +} + +static inline size_t get_size_in_bytes(Data_type dtype) { return get_size_in_bytes(1, dtype); } + +static inline size_t get_size_in_bits(Data_type dtype) { + switch (dtype) { + case DATA_TYPE_FP32: + return 32; + case DATA_TYPE_FP16: + return 16; + case DATA_TYPE_INT32: + return 32; + case DATA_TYPE_INT8: + return 8; + case DATA_TYPE_BF16: + return 16; + case DATA_TYPE_E2M1: + return 4; + case DATA_TYPE_E4M3: + return 8; + case DATA_TYPE_E5M2: + return 8; + default: + TORCH_CHECK(false, "FMHA Data Type is not supported."); + return 0; + } +} +constexpr int32_t kSM_70 = 70; +constexpr int32_t kSM_72 = 72; +constexpr int32_t kSM_75 = 75; +constexpr int32_t kSM_80 = 80; +constexpr int32_t kSM_86 = 86; +constexpr int32_t kSM_89 = 89; +constexpr int32_t kSM_90 = 90; +constexpr int32_t kSM_100 = 100; +constexpr int32_t kSM_120 = 120; diff --git a/include/flashinfer/trtllm/common/cudaBf16Fallbacks.cuh b/include/flashinfer/trtllm/common/cudaBf16Fallbacks.cuh new file mode 100644 index 0000000000000000000000000000000000000000..993038aa7fbd74a8f40645caec68775163f20e7e --- /dev/null +++ b/include/flashinfer/trtllm/common/cudaBf16Fallbacks.cuh @@ -0,0 +1,292 @@ +/* + * Copyright (c) 2019-2023, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include + +#include "flashinfer/trtllm/common/cudaBf16Wrapper.h" + +namespace tensorrt_llm { +namespace common { + +#ifdef ENABLE_BF16 +inline __device__ float2 bf1622float2(const __nv_bfloat162 val) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + float2 f_val; + f_val.x = __low2float(val); + f_val.y = __high2float(val); + return f_val; +#else + return __bfloat1622float2(val); +#endif +} + +inline __device__ int16_t bf1622int16(__nv_bfloat162 val) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + float2 f_val; + f_val.x = max(min(__low2float(val), 127.f), -128.f); + f_val.y = max(min(__high2float(val), 127.f), -128.f); + + union { + int8_t int8[2]; + int16_t int16; + }; + + int8[0] = static_cast(static_cast(f_val.x)); + int8[1] = static_cast(static_cast(f_val.y)); + return int16; +#else + val = __hmin2(val, make_bfloat162(127., 127.)); + val = __hmax2(val, make_bfloat162(-128., -128.)); + + union { + int8_t int8[2]; + int16_t int16; + }; + + int8[0] = static_cast(static_cast(val.x)); + int8[1] = static_cast(static_cast(val.y)); + return int16; +#endif +} + +inline __device__ __nv_bfloat162 float22bf162(const float2 val) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + return __floats2bfloat162_rn(val.x, val.y); +#else + return __float22bfloat162_rn(val); +#endif +} + +inline __device__ __nv_bfloat162 bf162bf162(const __nv_bfloat16 val) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + __nv_bfloat162 val2; + val2.x = val; + val2.y = val; + return val2; +#else + return __bfloat162bfloat162(val); +#endif +} + +inline __device__ __nv_bfloat162 bf16hadd2(const __nv_bfloat162 x, const __nv_bfloat162 y) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + float fxl, fxh, fyl, fyh; + fxl = __low2float(x); + fxh = __high2float(x); + fyl = __low2float(y); + fyh = __high2float(y); + return __floats2bfloat162_rn(fxl + fyl, fxh + fyh); +#else + return __hadd2(x, y); +#endif +} + +inline __device__ __nv_bfloat16 bf16hadd(const __nv_bfloat16 x, const __nv_bfloat16 y) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + return __float2bfloat16(__bfloat162float(x) + __bfloat162float(y)); +#else + return __hadd(x, y); +#endif +} + +inline __device__ __nv_bfloat162 bf16hsub2(const __nv_bfloat162 x, const __nv_bfloat162 y) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + float fxl, fxh, fyl, fyh; + fxl = __low2float(x); + fxh = __high2float(x); + fyl = __low2float(y); + fyh = __high2float(y); + return __floats2bfloat162_rn(fxl - fyl, fxh - fyh); +#else + return __hsub2(x, y); +#endif +} + +inline __device__ __nv_bfloat16 bf16hsub(const __nv_bfloat16 x, const __nv_bfloat16 y) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + return __float2bfloat16(__bfloat162float(x) - __bfloat162float(y)); +#else + return __hsub(x, y); +#endif +} + +inline __device__ __nv_bfloat162 bf16hmul2(const __nv_bfloat162 x, const __nv_bfloat162 y) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + float fxl, fxh, fyl, fyh; + fxl = __low2float(x); + fxh = __high2float(x); + fyl = __low2float(y); + fyh = __high2float(y); + return __floats2bfloat162_rn(fxl * fyl, fxh * fyh); +#else + return __hmul2(x, y); +#endif +} + +inline __device__ __nv_bfloat16 bf16hmul(const __nv_bfloat16 x, const __nv_bfloat16 y) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + return __float2bfloat16(__bfloat162float(x) * __bfloat162float(y)); +#else + return __hmul(x, y); +#endif +} + +inline __device__ __nv_bfloat162 bf16hfma2(const __nv_bfloat162 x, const __nv_bfloat162 y, + const __nv_bfloat162 z) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + float fxl, fxh, fyl, fyh, fzl, fzh; + fxl = __low2float(x); + fxh = __high2float(x); + fyl = __low2float(y); + fyh = __high2float(y); + fzl = __low2float(z); + fzh = __high2float(z); + return __floats2bfloat162_rn(fxl * fyl + fzl, fxh * fyh + fzh); +#else + return __hfma2(x, y, z); +#endif +} + +inline __device__ __nv_bfloat16 bf16hfma(const __nv_bfloat16 x, const __nv_bfloat16 y, + const __nv_bfloat16 z) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + return __float2bfloat16(__bfloat162float(x) * __bfloat162float(y) + __bfloat162float(z)); +#else + return __hfma(x, y, z); +#endif +} + +inline __device__ __nv_bfloat162 bf16exp2(const __nv_bfloat162 x) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + float fxl, fxh; + fxl = __low2float(x); + fxh = __high2float(x); + ; + return __floats2bfloat162_rn(expf(fxl), expf(fxh)); +#else + return h2exp(x); +#endif +} + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 800) +#if defined(CUDART_VERSION) && (CUDART_VERSION < 12020) + +inline __device__ __nv_bfloat162 make_bfloat162(const __nv_bfloat16 x, const __nv_bfloat16 y) { + __nv_bfloat162 t; + t.x = x; + t.y = y; + return t; +} +#endif +#endif + +inline __device__ __nv_bfloat16 bf16hadd(__nv_bfloat16 a, __nv_bfloat16 b, __nv_bfloat16 c) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + return __float2bfloat16(__bfloat162float(a) + __bfloat162float(b) + __bfloat162float(c)); +#else + return a + b + c; +#endif +} + +inline __device__ __nv_bfloat16 bf16hadd(__nv_bfloat16 a, __nv_bfloat16 b, __nv_bfloat16 c, + __nv_bfloat16 d) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + return __float2bfloat16(__bfloat162float(a) + __bfloat162float(b) + __bfloat162float(c) + + __bfloat162float(d)); +#else + return (__nv_bfloat16)((float)a + (float)b + (float)c + (float)d); +#endif +} + +inline __device__ __nv_bfloat162 bf16hadd2(__nv_bfloat162 a, __nv_bfloat162 b, __nv_bfloat162 c) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + float fal, fah, fbl, fbh, fcl, fch; + fal = __low2float(a); + fah = __high2float(a); + fbl = __low2float(b); + fbh = __high2float(b); + fcl = __low2float(c); + fch = __high2float(c); + return __floats2bfloat162_rn(fal + fbl + fcl, fah + fbh + fch); +#else + return a + b + c; +#endif +} + +inline __device__ __nv_bfloat16 bf16hmul(__nv_bfloat16 a, __nv_bfloat16 b, __nv_bfloat16 c) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + return __float2bfloat16(__bfloat162float(a) * __bfloat162float(b) * __bfloat162float(c)); +#else + return a * b * c; +#endif +} + +inline __device__ __nv_bfloat162 bf16hmul2(__nv_bfloat162 a, __nv_bfloat162 b, __nv_bfloat162 c) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + float fal, fah, fbl, fbh, fcl, fch; + fal = __low2float(a); + fah = __high2float(a); + fbl = __low2float(b); + fbh = __high2float(b); + fcl = __low2float(c); + fch = __high2float(c); + return __floats2bfloat162_rn(fal * fbl * fcl, fah * fbh * fch); +#else + return a * b * c; +#endif +} + +inline __device__ __nv_bfloat162 bf16hfma2(__nv_bfloat162 a, __nv_bfloat162 b, __nv_bfloat162 c, + __nv_bfloat162 d) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + float fal, fah, fbl, fbh, fcl, fch, fdl, fdh; + fal = __low2float(a); + fah = __high2float(a); + fbl = __low2float(b); + fbh = __high2float(b); + fcl = __low2float(c); + fch = __high2float(c); + fdl = __low2float(d); + fdh = __high2float(d); + return __floats2bfloat162_rn(fal * fbl * fcl + fdl, fah * fbh * fch + fdh); +#else + return a * b * c + d; +#endif +} + +#endif // ENABLE_BF16 + +} // namespace common +} // namespace tensorrt_llm + +// Operator definitions intentionally in global namespace +namespace { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 800) +#if defined(CUDART_VERSION) && (CUDART_VERSION < 12020) + +inline __device__ __nv_bfloat162 operator*(const __nv_bfloat162 x, const __nv_bfloat162 y) { + return tensorrt_llm::common::bf16hmul2(x, y); +}; + +inline __device__ __nv_bfloat162 operator+(const __nv_bfloat162 x, const __nv_bfloat162 y) { + return tensorrt_llm::common::bf16hadd2(x, y); +}; +#endif +#endif +} // namespace diff --git a/include/flashinfer/trtllm/common/cudaBf16Wrapper.h b/include/flashinfer/trtllm/common/cudaBf16Wrapper.h new file mode 100644 index 0000000000000000000000000000000000000000..fb2a89af5cd609bc1231ec28bf9f567d6c1e4b7a --- /dev/null +++ b/include/flashinfer/trtllm/common/cudaBf16Wrapper.h @@ -0,0 +1,21 @@ +/* + * Copyright (c) 2019-2023, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#ifdef ENABLE_BF16 +#include +#endif diff --git a/include/flashinfer/trtllm/common/cudaFp8Utils.h b/include/flashinfer/trtllm/common/cudaFp8Utils.h new file mode 100644 index 0000000000000000000000000000000000000000..02efb1c995a0d905d86aa7a8de85da5f69dd4f1b --- /dev/null +++ b/include/flashinfer/trtllm/common/cudaFp8Utils.h @@ -0,0 +1,281 @@ +/* + * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#ifdef ENABLE_FP8 +#include +#include +#include + +#define FP8_MHA +#define FUSE_GEMM_ACT +#define FP8_GEMM_OUTPUT_QUANT_DISABLE + +#ifdef FUSE_GEMM_ACT +#define USE_QGMMA +#endif + +namespace tensorrt_llm { +namespace common { + +constexpr float FP8_E4M3_MAX = 448.0f; + +enum QuantizeMode { + PER_CHANNEL, + PER_TENSOR, + PER_CHANNEL_WEIGHT_PER_TENSOR_ACT, + PER_TOKEN, +}; + +// Packed Data Type +typedef struct __CUDA_ALIGN__(32) { + float array[8]; +} float8; + +typedef struct __CUDA_ALIGN__(16) { + half array[8]; +} half8; + +typedef struct __CUDA_ALIGN__(8) { + half2 array[2]; +} half2_2; + +typedef struct __CUDA_ALIGN__(8) { + half array[4]; +} half_4; + +#ifdef ENABLE_BF16 +typedef struct __CUDA_ALIGN__(4) { + __nv_bfloat16 array[2]; +} __nv_bfloat16_2; + +typedef struct __CUDA_ALIGN__(8) { + __nv_bfloat162 x, y; +} __nv_bfloat162_2_xy; + +typedef struct __CUDA_ALIGN__(8) { + __nv_bfloat16 array[4]; +} __nv_bfloat164; + +typedef struct __CUDA_ALIGN__(8) { + __nv_bfloat162 array[2]; +} __nv_bfloat162_2; + +typedef struct __CUDA_ALIGN__(16) { + __nv_bfloat16 array[8]; +} __nv_bfloat168; + +typedef struct __CUDA_ALIGN__(16) { + __nv_bfloat162 array[4]; +} __nv_bfloat162_4; + +typedef struct __CUDA_ALIGN__(32) { + __nv_bfloat16 array[16]; +} __nv_bfloat1616; +#endif + +#ifdef ENABLE_FP8 +typedef struct __CUDA_ALIGN__(2) { + __nv_fp8_e4m3 array[2]; +} __nv_fp8_2_e4m3; + +typedef struct __CUDA_ALIGN__(4) { + __nv_fp8_e4m3 array[4]; +} __nv_fp8_4_e4m3; + +typedef struct __CUDA_ALIGN__(4) { + __nv_fp8x2_e4m3 array[2]; +} __nv_fp8x2_x2_e4m3; + +typedef struct __CUDA_ALIGN__(8) { + __nv_fp8_e4m3 array[8]; +} __nv_fp8_8_e4m3; + +typedef struct __CUDA_ALIGN__(8) { + __nv_fp8x2_e4m3 array[4]; +} __nv_fp8x2_x4_e4m3; + +typedef struct __CUDA_ALIGN__(16) { + __nv_fp8_e4m3 array[16]; +} __nv_fp8x16_e4m3; +#endif + +// only BF16 and FP8 +template +struct PackType { + using type = float; +}; + +template +struct PackType { + struct __CUDA_ALIGN__(std::alignment_of_v) type { + T array[1]; + }; +}; + +template <> +struct PackType { + struct __CUDA_ALIGN__(8) type { + float array[2]; + }; +}; + +template <> +struct PackType { + struct __CUDA_ALIGN__(16) type { + float array[4]; + }; +}; + +template <> +struct PackType { + struct __CUDA_ALIGN__(32) type { + float array[8]; + }; +}; + +template <> +struct PackType { + struct __CUDA_ALIGN__(64) type { + float array[16]; + }; +}; + +template <> +struct PackType { + struct __CUDA_ALIGN__(4) type { + half array[2]; + }; +}; + +template <> +struct PackType { + struct __CUDA_ALIGN__(8) type { + half array[4]; + }; +}; + +template <> +struct PackType { + struct __CUDA_ALIGN__(16) type { + half array[8]; + }; +}; + +#ifdef ENABLE_BF16 +template <> +struct PackType<__nv_bfloat16, 2> { + using type = __nv_bfloat16_2; +}; + +template <> +struct PackType<__nv_bfloat16, 4> { + using type = __nv_bfloat164; +}; + +template <> +struct PackType<__nv_bfloat16, 8> { + using type = __nv_bfloat168; +}; + +template <> +struct PackType<__nv_bfloat16, 16> { + using type = __nv_bfloat1616; +}; +#endif + +#ifdef ENABLE_FP8 +template <> +struct PackType<__nv_fp8_e4m3, 2> { + using type = __nv_fp8_2_e4m3; +}; + +template <> +struct PackType<__nv_fp8_e4m3, 4> { + using type = __nv_fp8_4_e4m3; +}; + +template <> +struct PackType<__nv_fp8_e4m3, 8> { + using type = __nv_fp8_8_e4m3; +}; + +template <> +struct PackType<__nv_fp8_e4m3, 16> { + using type = __nv_fp8x16_e4m3; +}; + +#endif + +__inline__ __device__ void fp8x4_e4m3_to_bfloat2(__nv_bfloat162* out1, __nv_bfloat162* out2, + __nv_fp8x4_e4m3 const* in) { + const char4 tmp_val = reinterpret_cast(in)[0]; + *out1 = __nv_bfloat162((float)reinterpret_cast<__nv_fp8_e4m3 const*>(&tmp_val.x)[0], + (float)reinterpret_cast<__nv_fp8_e4m3 const*>(&tmp_val.y)[0]); + *out2 = __nv_bfloat162((float)reinterpret_cast<__nv_fp8_e4m3 const*>(&tmp_val.z)[0], + (float)reinterpret_cast<__nv_fp8_e4m3 const*>(&tmp_val.w)[0]); +} + +__inline__ __device__ __nv_bfloat162 fp8x2_e4m3_to_bfloat2(__nv_fp8x2_e4m3 const* in) { + const char2 tmp_val = reinterpret_cast(in)[0]; + __nv_bfloat162 out = __nv_bfloat162((float)reinterpret_cast<__nv_fp8_e4m3 const*>(&tmp_val.x)[0], + (float)reinterpret_cast<__nv_fp8_e4m3 const*>(&tmp_val.y)[0]); + return out; +} + +__inline__ __device__ void fp8x4_e4m3_to_half2(half2* out1, half2* out2, + __nv_fp8x4_e4m3 const* in) { + const char4 tmp_val = reinterpret_cast(in)[0]; + *out1 = half2((float)reinterpret_cast<__nv_fp8_e4m3 const*>(&tmp_val.x)[0], + (float)reinterpret_cast<__nv_fp8_e4m3 const*>(&tmp_val.y)[0]); + *out2 = half2((float)reinterpret_cast<__nv_fp8_e4m3 const*>(&tmp_val.z)[0], + (float)reinterpret_cast<__nv_fp8_e4m3 const*>(&tmp_val.w)[0]); +} + +__inline__ __device__ half2 fp8x2_e4m3_to_half2(__nv_fp8x2_e4m3 const* in) { + const char2 tmp_val = reinterpret_cast(in)[0]; + half2 out = half2((float)reinterpret_cast<__nv_fp8_e4m3 const*>(&tmp_val.x)[0], + (float)reinterpret_cast<__nv_fp8_e4m3 const*>(&tmp_val.y)[0]); + return out; +} + +template +void invokeQuantizeMatrix(T_OUT* output, T_S const* input_qua_amax_ptr, T_IN const* input, + int64_t numel, int64_t lda, QuantizeMode quantize_mode, + cudaStream_t stream); + +template +void invokeDequantizeMatrix(T_OUT* output, T_S const* input_qua_amax_ptr, T_IN const* input, + int64_t numel, int64_t lda, QuantizeMode quantize_mode, + cudaStream_t stream); + +template +void invokeFakeQuantize(T_OUT* dst, const T_IN* src, const int64_t numel, cudaStream_t stream); + +template +void invokeComputeFP8QuantizeScale(T_S* quant_ptr, const T_W* weights, const int64_t k, + const int64_t lda, QuantizeMode quantize_mode, + cudaStream_t stream); + +template +void invokeComputeScalesAndQuantizeMatrix(T_OUT* output, T_S* quant_ptr, const T_IN* weights, + const int64_t numel, const int64_t lda, + QuantizeMode quantize_mode, cudaStream_t stream); + +} // namespace common +} // namespace tensorrt_llm +#endif // ENABLE_FP8 diff --git a/include/flashinfer/trtllm/common/cudaTypeUtils.cuh b/include/flashinfer/trtllm/common/cudaTypeUtils.cuh new file mode 100644 index 0000000000000000000000000000000000000000..456af48e01dd96f24860fc065b206bfcc239c7b1 --- /dev/null +++ b/include/flashinfer/trtllm/common/cudaTypeUtils.cuh @@ -0,0 +1,647 @@ +/* + * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include + +#include "flashinfer/trtllm/common/cudaBf16Fallbacks.cuh" +#include "flashinfer/trtllm/common/cudaBf16Wrapper.h" +#include "flashinfer/trtllm/common/cudaFp8Utils.h" +#if ENABLE_BF16 +#include +#endif + +namespace tensorrt_llm { +namespace common { + +template +inline __device__ T ldg(T const* val) { + return __ldg(val); +} + +#if ENABLE_BF16 +template <> +inline __device__ __nv_bfloat162 ldg(__nv_bfloat162 const* val) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + return val[0]; +#else + return __ldg(val); +#endif +} + +template <> +inline __device__ __nv_bfloat16 ldg(__nv_bfloat16 const* val) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + return val[0]; +#else + return __ldg(val); +#endif +} +#endif // ENABLE_BF16 + +// Get type2 from type or vice versa (applied to half and bfloat16) +template +struct TypeConverter { + using Type = half2; +}; // keep for generality + +template <> +struct TypeConverter { + using Type = half; +}; + +template <> +struct TypeConverter { + using Type = half2; +}; + +#if ENABLE_BF16 +template <> +struct TypeConverter<__nv_bfloat162> { + using Type = __nv_bfloat16; +}; + +template <> +struct TypeConverter<__nv_bfloat16> { + using Type = __nv_bfloat162; +}; +#endif // ENABLE_BF16 + +// Defined math operations (bfloat16 fallback to fp32 when it is not supported) +template +inline __device__ T hadd2(T a, T b) { + return __hadd2(a, b); +} + +#if ENABLE_BF16 +template <> +inline __device__ __nv_bfloat162 hadd2(__nv_bfloat162 a, __nv_bfloat162 b) { + return bf16hadd2(a, b); +} +#endif // ENABLE_BF16 + +template +inline __device__ T add(T a, T b) { + return a + b; +} + +template <> +inline __device__ half2 add(half2 a, half2 b) { + return __hadd2(a, b); +} + +template <> +inline __device__ half add(half a, half b) { + return __hadd(a, b); +} + +#if ENABLE_BF16 +template <> +inline __device__ __nv_bfloat162 add(__nv_bfloat162 a, __nv_bfloat162 b) { + return bf16hadd2(a, b); +} + +template <> +inline __device__ __nv_bfloat16 add(__nv_bfloat16 a, __nv_bfloat16 b) { + return bf16hadd(a, b); +} + +inline __device__ __nv_bfloat16 add(__nv_bfloat16 a, float b) { + return bf16hadd(a, __float2bfloat16(b)); +} +#endif // ENABLE_BF16 + +// applies to all 4 values addition +template +inline __device__ T add(T a, T b, T c) { + return a + b + c; +} + +#if ENABLE_BF16 +inline __device__ __nv_bfloat16 add(__nv_bfloat16 a, __nv_bfloat16 b, __nv_bfloat16 c) { + return bf16hadd(a, b, c); +} + +inline __device__ __nv_bfloat162 add(__nv_bfloat162 a, __nv_bfloat162 b, __nv_bfloat162 c) { + return bf16hadd2(a, b, c); +} +#endif // ENABLE_BF16 + +// applies to all 4 values addition +template +inline __device__ T add(T a, T b, T c, T d) { + return (T)((float)a + (float)b + (float)c + (float)d); +} + +#if ENABLE_BF16 +inline __device__ __nv_bfloat16 add(__nv_bfloat16 a, __nv_bfloat16 b, __nv_bfloat16 c, + __nv_bfloat16 d) { + return bf16hadd(a, b, c, d); +} +#endif // ENABLE_BF16 + +template +inline __device__ T hsub2(T a, T b) { + return __hsub2(a, b); +} + +#if ENABLE_BF16 +template <> +inline __device__ __nv_bfloat162 hsub2(__nv_bfloat162 a, __nv_bfloat162 b) { + return bf16hsub2(a, b); +} +#endif // ENABLE_BF16 + +template +inline __device__ T hmul2(T a, T b) { + return __hmul2(a, b); +} + +#if ENABLE_BF16 +template <> +inline __device__ __nv_bfloat162 hmul2(__nv_bfloat162 a, __nv_bfloat162 b) { + return bf16hmul2(a, b); +} +#endif // ENABLE_BF16 + +template +inline __device__ T hmul2(T a, T b, T c) { + return a * b * c; +} + +#if ENABLE_BF16 +template <> +inline __device__ __nv_bfloat162 hmul2(__nv_bfloat162 a, __nv_bfloat162 b, __nv_bfloat162 c) { + return bf16hmul2(a, b, c); +} +#endif // ENABLE_BF16 + +template +inline __device__ T mul(T a, T b, T c) { + return a * b * c; +} + +#if ENABLE_BF16 +template <> +inline __device__ __nv_bfloat16 mul(__nv_bfloat16 a, __nv_bfloat16 b, __nv_bfloat16 c) { + return bf16hmul(a, b, c); +} + +inline __device__ __nv_bfloat162 mul(__nv_bfloat162 a, __nv_bfloat162 b, __nv_bfloat162 c) { + return bf16hmul2(a, b, c); +} +#endif // ENABLE_BF16 + +template +inline __device__ T fma(T a, T b, T c, T d) { + return a * b * c + d; +} + +#if ENABLE_BF16 +inline __device__ __nv_bfloat162 fma(__nv_bfloat162 a, __nv_bfloat162 b, __nv_bfloat162 c, + __nv_bfloat162 d) { + return bf16hfma2(a, b, c, d); +} +#endif // ENABLE_BF16 + +template +inline __device__ T fma(T a, T b, T c) { + return a * b + c; +} + +#if ENABLE_BF16 +template <> +inline __device__ __nv_bfloat162 fma(__nv_bfloat162 a, __nv_bfloat162 b, __nv_bfloat162 c) { + return bf16hfma2(a, b, c); +} + +template <> +inline __device__ __nv_bfloat16 fma(__nv_bfloat16 a, __nv_bfloat16 b, __nv_bfloat16 c) { + return bf16hfma(a, b, c); +} +#endif // ENABLE_BF16 + +template +inline __device__ T hexp2(T a) { + return h2exp(a); +} + +#if ENABLE_BF16 +template <> +inline __device__ __nv_bfloat162 hexp2(__nv_bfloat162 a) { + return bf16exp2(a); +} +#endif // ENABLE_BF16 + +template +__device__ inline T_OUT cuda_cast(T_IN val) { + return val; +} + +template <> +__device__ inline float2 cuda_cast(int2 val) { + return make_float2(val.x, val.y); +} + +template <> +__device__ inline float2 cuda_cast(float val) { + return make_float2(val, val); +} + +template <> +__device__ inline float2 cuda_cast(half2 val) { + return __half22float2(val); +} + +template <> +__device__ inline half2 cuda_cast(float2 val) { + return __float22half2_rn(val); +} + +template <> +__device__ inline half2 cuda_cast(float val) { + return __float2half2_rn(val); +} + +template <> +__device__ inline half2 cuda_cast(half val) { + return __half2half2(val); +} + +template <> +__device__ inline int8_t cuda_cast(half val) { + union { + int8_t int8[2]; + int16_t int16; + }; + + union { + half fp16; + int16_t int16_in; + }; + + fp16 = val; + asm volatile("cvt.rni.sat.s8.f16 %0, %1;" : "=h"(int16) : "h"(int16_in)); + return int8[0]; +} + +template <> +__device__ inline int16_t cuda_cast(half2 val) { + union { + int8_t int8[2]; + int16_t int16; + }; + + int8[0] = cuda_cast(val.x); + int8[1] = cuda_cast(val.y); + return int16; +} + +template <> +__device__ inline int8_t cuda_cast(float val) { + union { + int8_t int8[2]; + int16_t int16; + }; + + asm volatile("cvt.rni.sat.s8.f32 %0, %1;" : "=h"(int16) : "f"(val)); + return int8[0]; +} + +template <> +__device__ inline int16_t cuda_cast(float2 val) { + union { + int8_t int8[2]; + int16_t int16; + }; + + int8[0] = cuda_cast(val.x); + int8[1] = cuda_cast(val.y); + return int16; +} + +template <> +__device__ inline half2 cuda_cast(int16_t val) { + union { + int8_t int8[2]; + int16_t int16; + }; + + int16 = val; + return make_half2(int8[0], int8[1]); +} + +template <> +__device__ inline float2 cuda_cast(int16_t val) { + union { + int8_t int8[2]; + int16_t int16; + }; + + int16 = val; + return make_float2(int8[0], int8[1]); +} + +#ifdef ENABLE_BF16 +template <> +__device__ inline __nv_bfloat16 cuda_cast(int32_t val) { + return static_cast(val); +} + +template <> +__device__ inline __nv_bfloat16 cuda_cast(int8_t val) { + return static_cast(val); +} + +template <> +__device__ inline int8_t cuda_cast(__nv_bfloat16 val) { + return static_cast(val); +} + +template <> +__device__ inline float cuda_cast(__nv_bfloat16 val) { + return __bfloat162float(val); +} + +template <> +__device__ inline float2 cuda_cast(__nv_bfloat162 val) { + return bf1622float2(val); +} + +template <> +__device__ inline half cuda_cast(__nv_bfloat16 val) { + return __float2half(__bfloat162float(val)); +} + +template <> +__device__ inline int16_t cuda_cast(__nv_bfloat162 val) { + return bf1622int16(val); +} + +template <> +__device__ inline __nv_bfloat16 cuda_cast<__nv_bfloat16, float>(float val) { + return __float2bfloat16(val); +} + +template <> +__device__ inline __nv_bfloat16 cuda_cast<__nv_bfloat16, half>(half val) { + return __float2bfloat16(__half2float(val)); +} + +template <> +__device__ inline __nv_bfloat162 cuda_cast<__nv_bfloat162, __nv_bfloat16>(__nv_bfloat16 val) { + return bf162bf162(val); +} + +template <> +__device__ inline __nv_bfloat162 cuda_cast<__nv_bfloat162, float>(float val) { + return __float2bfloat162_rn(val); +} + +template <> +__device__ inline __nv_bfloat162 cuda_cast<__nv_bfloat162, float2>(float2 val) { + return float22bf162(val); +} + +template <> +__device__ inline __nv_bfloat162 cuda_cast<__nv_bfloat162, int16_t>(int16_t val) { + union { + int8_t int8[2]; + int16_t int16; + }; + + int16 = val; + __nv_bfloat162 res; + res.x = cuda_cast<__nv_bfloat16>(int8[0]); + res.y = cuda_cast<__nv_bfloat16>(int8[1]); + return res; +} + +template <> +__device__ inline __nv_bfloat162 cuda_cast<__nv_bfloat162, half2>(half2 val) { + return float22bf162(__half22float2(val)); +} + +#endif // ENABLE BF16 + +template +__device__ inline T cuda_abs(T val) { + assert(false); + return {}; +} + +template <> +__device__ inline float cuda_abs(float val) { + return fabs(val); +} + +template <> +__device__ inline float2 cuda_abs(float2 val) { + return make_float2(fabs(val.x), fabs(val.y)); +} + +template <> +__device__ inline half cuda_abs(half val) { + return __habs(val); +} + +template <> +__device__ inline half2 cuda_abs(half2 val) { + return __habs2(val); +} + +#ifdef ENABLE_BF16 + +#if __CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__) +template <> +__device__ inline __nv_bfloat16 cuda_abs(__nv_bfloat16 val) { + return __habs(val); +} + +template <> +__device__ inline __nv_bfloat162 cuda_abs(__nv_bfloat162 val) { + return __habs2(val); +} +#endif + +#endif // ENABLE_FP16 + +template +__device__ inline To cuda_sum(Ti val) { + return cuda_cast(val); +}; + +template +__device__ inline To cuda_sum(float2 val) { + return cuda_cast(val.x + val.y); +}; + +// Unary maximum: compute the max of a vector type +template +__device__ inline To cuda_max(Ti val) { + return cuda_cast(val); +}; + +template <> +__device__ inline float cuda_max(float2 val) { + return fmaxf(val.x, val.y); +} + +template <> +__device__ inline half cuda_max(half2 val) { + return __hmax(val.x, val.y); +} + +#ifdef ENABLE_BF16 +template <> +__device__ inline __nv_bfloat16 cuda_max(__nv_bfloat162 val) { +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)) + return __hmax(val.x, val.y); +#else + assert(0); + asm volatile("brkpt;\n" ::); + return __nv_bfloat16(0); +#endif +} +#endif + +// Binary maximum: compute the max of two values. +template +__device__ inline T cuda_max(T val1, T val2) { + return (val1 > val2) ? val1 : val2; +} + +template <> +__device__ inline float2 cuda_max(float2 val1, float2 val2) { + float2 out; + out.x = fmaxf(val1.x, val2.x); + out.y = fmaxf(val1.y, val2.y); + return out; +} + +template <> +__device__ inline half2 cuda_max(half2 val1, half2 val2) { + return __hmax2(val1, val2); +} + +#ifdef ENABLE_BF16 +template <> +__device__ inline __nv_bfloat162 cuda_max(__nv_bfloat162 val1, __nv_bfloat162 val2) { + return __hmax2(val1, val2); +} +#endif // ENABLE_BF16 + +// Binary maximum: compute the min of two values. +template +__device__ inline T cuda_min(T val1, T val2) { + return (val1 < val2) ? val1 : val2; +} + +template <> +__device__ inline float2 cuda_min(float2 val1, float2 val2) { + float2 out; + out.x = fminf(val1.x, val2.x); + out.y = fminf(val1.y, val2.y); + return out; +} + +template <> +__device__ inline half2 cuda_min(half2 val1, half2 val2) { + return __hmin2(val1, val2); +} + +#ifdef ENABLE_BF16 +template <> +__device__ inline __nv_bfloat162 cuda_min(__nv_bfloat162 val1, __nv_bfloat162 val2) { + return __hmin2(val1, val2); +} +#endif // ENABLE_BF16 + +// Helper function of clamping the val into the given range. +template +inline __device__ T cuda_clamp(T val, T minVal, T maxVal) { + return cuda_min(cuda_max(val, minVal), maxVal); +} + +#ifdef ENABLE_FP8 +template <> +__device__ inline float2 cuda_cast(__nv_fp8x2_e4m3 val) { + return bf1622float2(fp8x2_e4m3_to_bfloat2(&val)); +} + +template <> +__device__ inline half2 cuda_cast(__nv_fp8x2_e4m3 val) { + return fp8x2_e4m3_to_half2(&val); +} + +template <> +__device__ inline __nv_fp8x2_e4m3 cuda_cast<__nv_fp8x2_e4m3, float2>(float2 val) { + return __nv_fp8x2_e4m3(bf1622float2(float22bf162(val))); +} + +template <> +__device__ inline __nv_fp8x2_e4m3 cuda_cast<__nv_fp8x2_e4m3, half2>(half2 val) { + return __nv_fp8x2_e4m3(cuda_cast(val)); +} + +template <> +__device__ inline __nv_fp8x2_e4m3 cuda_cast<__nv_fp8x2_e4m3, __nv_bfloat162>(__nv_bfloat162 val) { + return __nv_fp8x2_e4m3(cuda_cast(val)); +} + +template <> +__device__ inline __nv_fp8_e4m3 cuda_cast<__nv_fp8_e4m3, half>(half val) { + return __nv_fp8_e4m3(val); +} + +template <> +__device__ inline __nv_fp8_e4m3 cuda_cast<__nv_fp8_e4m3, __nv_bfloat16>(__nv_bfloat16 val) { + return __nv_fp8_e4m3(val); +} + +template <> +__device__ inline __nv_fp8_e4m3 cuda_cast<__nv_fp8_e4m3, float>(float val) { + return __nv_fp8_e4m3(val); +} + +template <> +__device__ inline float cuda_cast(__nv_fp8_e4m3 val) { + return (float)val; +} + +template <> +__device__ inline __nv_bfloat162 cuda_cast<__nv_bfloat162, __nv_fp8x2_e4m3>(__nv_fp8x2_e4m3 val) { + return fp8x2_e4m3_to_bfloat2(&val); +} + +template <> +__device__ inline int8_t cuda_cast(__nv_fp8_e4m3 val) { + // no impl + return 0; +} + +template <> +__device__ inline __nv_fp8_e4m3 cuda_cast<__nv_fp8_e4m3, int8_t>(int8_t val) { + return cuda_cast<__nv_fp8_e4m3>(cuda_cast<__nv_bfloat16>(cuda_cast(val))); +} + +#endif // ENABLE_FP8 + +} // namespace common +} // namespace tensorrt_llm diff --git a/include/flashinfer/trtllm/common/cudaUtils.h b/include/flashinfer/trtllm/common/cudaUtils.h new file mode 100644 index 0000000000000000000000000000000000000000..2fa1a4d0c47c01ca9c17d3f08802dce3b2dd4a05 --- /dev/null +++ b/include/flashinfer/trtllm/common/cudaUtils.h @@ -0,0 +1,158 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights + * reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +// #include "tensorrt_llm/common/cudaBf16Wrapper.h" +// #include "tensorrt_llm/common/cudaDriverWrapper.h" +// #include "tensorrt_llm/common/cudaFp8Utils.h" +// #include "tensorrt_llm/common/logger.h" +// #include "tensorrt_llm/common/tllmException.h" + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +// #ifndef _WIN32 // Linux +// #include +// #endif // not WIN32 +// #include +// #ifdef _WIN32 // Windows +// #include +// #undef ERROR // A Windows header file defines ERROR as 0, but it's used in our logger.h enum. +// Logging breaks without +// // this undef. +// #endif // WIN32 + +namespace tensorrt_llm::common { + +// // workspace for cublas gemm : 32MB +// #define CUBLAS_WORKSPACE_SIZE 33554432 + +// typedef struct __align__(4) +// { +// half x, y, z, w; +// } + +// half4; + +// /* **************************** type definition ***************************** */ + +// enum CublasDataType +// { +// FLOAT_DATATYPE = 0, +// HALF_DATATYPE = 1, +// BFLOAT16_DATATYPE = 2, +// INT8_DATATYPE = 3, +// FP8_DATATYPE = 4 +// }; + +// enum TRTLLMCudaDataType +// { +// FP32 = 0, +// FP16 = 1, +// BF16 = 2, +// INT8 = 3, +// FP8 = 4 +// }; + +// enum class OperationType +// { +// FP32, +// FP16, +// BF16, +// INT8, +// FP8 +// }; + +/* **************************** debug tools ********************************* */ + +inline std::optional isCudaLaunchBlocking() { + thread_local bool firstCall = true; + thread_local std::optional result = std::nullopt; + if (!firstCall) { + char const* env = std::getenv("CUDA_LAUNCH_BLOCKING"); + if (env != nullptr && std::string(env) == "1") { + result = true; + } else { + result = false; + } + firstCall = false; + } + return result; +} + +inline std::optional isCapturing(cudaStream_t stream) { + cudaStreamCaptureStatus status; + TORCH_CHECK(cudaStreamIsCapturing(stream, &status) == cudaSuccess, + "CUDA error in cudaStreamIsCapturing"); + return status == cudaStreamCaptureStatus::cudaStreamCaptureStatusActive; +} + +inline bool doCheckError(cudaStream_t stream) { + auto const cudaLaunchBlocking = isCudaLaunchBlocking(); + if (cudaLaunchBlocking.has_value() && cudaLaunchBlocking.value()) { + return !isCapturing(stream); + } + +#ifndef NDEBUG + // Debug builds will sync when we're not capturing unless explicitly + // disabled. + bool const checkError = cudaLaunchBlocking.value_or(!isCapturing(stream)); +#else + bool const checkError = cudaLaunchBlocking.value_or(false); +#endif + + return checkError; +} + +inline void syncAndCheck(cudaStream_t stream, char const* const file, int const line) { + if (doCheckError(stream)) { + cudaStreamSynchronize(stream); + auto error = cudaGetLastError(); + TORCH_CHECK(error == cudaSuccess, "CUDA error in %s: %s", file, cudaGetErrorString(error)); + } +} + +#define sync_check_cuda_error(stream) tensorrt_llm::common::syncAndCheck(stream, __FILE__, __LINE__) + +template +inline size_t divUp(T1 const& a, T2 const& b) { + auto const tmp_a = static_cast(a); + auto const tmp_b = static_cast(b); + return (tmp_a + tmp_b - 1) / tmp_b; +} + +inline int roundUp(int a, int b) { return divUp(a, b) * b; } + +template ::value>, + typename = std::enable_if_t::value>> +auto constexpr ceilDiv(T numerator, U denominator) { + return (numerator + denominator - 1) / denominator; +} + +} // namespace tensorrt_llm::common diff --git a/include/flashinfer/trtllm/fmha/cubin/kernelMetaInfo.h b/include/flashinfer/trtllm/fmha/cubin/kernelMetaInfo.h new file mode 100644 index 0000000000000000000000000000000000000000..f69b2da15edb1727075acd2124094df85fda068a --- /dev/null +++ b/include/flashinfer/trtllm/fmha/cubin/kernelMetaInfo.h @@ -0,0 +1,170 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & + * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include +// Helper to print Data_type +inline const char* dataTypeToString(Data_type dt) { + switch (dt) { + case DATA_TYPE_FP16: + return "FP16"; + case DATA_TYPE_BF16: + return "BF16"; + case DATA_TYPE_FP32: + return "FP32"; + case DATA_TYPE_E4M3: + return "E4M3"; + case DATA_TYPE_E2M1: + return "E2M1"; + default: + return "UNKNOWN"; + } +} + +inline Data_type stringToDataType(std::string str) { + if (str == "DATA_TYPE_FP16") return DATA_TYPE_FP16; + if (str == "DATA_TYPE_BF16") return DATA_TYPE_BF16; + if (str == "DATA_TYPE_FP32") return DATA_TYPE_FP32; + if (str == "DATA_TYPE_E4M3") return DATA_TYPE_E4M3; + if (str == "DATA_TYPE_E2M1") return DATA_TYPE_E2M1; + return DATA_TYPE_UNKNOWN; +} + +inline int stringToArch(std::string str) { + if (str == "kSM_90") return kSM_90; + if (str == "kSM_100") return kSM_100; + if (str == "kSM_120") return kSM_120; + return 0; +} + +struct TllmGenFmhaKernelMetaInfo { + Data_type mDataTypeQ; + Data_type mDataTypeKv; + Data_type mDataTypeO; + int mTileSizeQ; + int mTileSizeKv; + int mStepQ; + int mStepKv; + int mHeadDimPerCtaV; + int mHeadDimQk; + int mHeadDimV; + int mSM; + const char* mFuncName; + int mSharedMemBytes; + int mThreadsPerCTA; + int mQkvLayout; + int mNumTokensPerPage; + int mMaskType; + int mKernelType; + int mMaxNumHeadsQPerKvInCta; + int mTileScheduler; + int mMultiCtasKvMode; + bool mGroupsHeadsQ; + bool mReuseSmemKForV; + bool m2CtaMma; + const char* sha256; + + void print() const { + std::cout << "TllmGenFmhaKernelMetaInfo {\n"; + std::cout << " mDataTypeQ: " << dataTypeToString(mDataTypeQ) << "\n"; + std::cout << " mDataTypeKv: " << dataTypeToString(mDataTypeKv) << "\n"; + std::cout << " mDataTypeO: " << dataTypeToString(mDataTypeO) << "\n"; + std::cout << " mTileSizeQ: " << mTileSizeQ << "\n"; + std::cout << " mTileSizeKv: " << mTileSizeKv << "\n"; + std::cout << " mStepQ: " << mStepQ << "\n"; + std::cout << " mStepKv: " << mStepKv << "\n"; + std::cout << " mHeadDimPerCtaV: " << mHeadDimPerCtaV << "\n"; + std::cout << " mHeadDimQk: " << mHeadDimQk << "\n"; + std::cout << " mHeadDimV: " << mHeadDimV << "\n"; + std::cout << " mSM: " << mSM << "\n"; + std::cout << " mFuncName: " << (mFuncName ? mFuncName : "null") << "\n"; + std::cout << " mSharedMemBytes: " << mSharedMemBytes << "\n"; + std::cout << " mThreadsPerCTA: " << mThreadsPerCTA << "\n"; + std::cout << " mQkvLayout: " << mQkvLayout << "\n"; + std::cout << " mNumTokensPerPage: " << mNumTokensPerPage << "\n"; + std::cout << " mMaskType: " << mMaskType << "\n"; + std::cout << " mKernelType: " << mKernelType << "\n"; + std::cout << " mMaxNumHeadsQPerKvInCta: " << mMaxNumHeadsQPerKvInCta << "\n"; + std::cout << " mTileScheduler: " << mTileScheduler << "\n"; + std::cout << " mGroupsHeadsQ: " << std::boolalpha << mGroupsHeadsQ << "\n"; + std::cout << " mMultiCtasKvMode: " << std::boolalpha << mMultiCtasKvMode << "\n"; + std::cout << " mReuseSmemKForV: " << std::boolalpha << mReuseSmemKForV << "\n"; + std::cout << " m2CtaMma: " << std::boolalpha << m2CtaMma << "\n"; + std::cout << " sha256: " << (sha256 ? sha256 : "null") << "\n"; + std::cout << "}\n"; + } + + static TllmGenFmhaKernelMetaInfo fromString(std::string code) { + std::vector param_list; + std::string current_param = ""; + for (int i = 0; i < code.size(); ++i) { + if (code[i] != ' ' && code[i] != ',' && code[i] != '"') { + current_param += code[i]; + } + if (code[i] == ',') { + param_list.push_back(current_param); + current_param = ""; + } + } + param_list.push_back(current_param); + assert(param_list.size() == 25); + const char* mFuncName = strdup(param_list[11].c_str()); + const char* sha256 = strdup(param_list[24].c_str()); + return TllmGenFmhaKernelMetaInfo{stringToDataType(param_list[0]), + stringToDataType(param_list[1]), + stringToDataType(param_list[2]), + std::stoi(param_list[3]), + std::stoi(param_list[4]), + std::stoi(param_list[5]), + std::stoi(param_list[6]), + std::stoi(param_list[7]), + std::stoi(param_list[8]), + std::stoi(param_list[9]), + stringToArch(param_list[10]), + mFuncName, + std::stoi(param_list[12]), + std::stoi(param_list[13]), + std::stoi(param_list[14]), + std::stoi(param_list[15]), + std::stoi(param_list[16]), + std::stoi(param_list[17]), + std::stoi(param_list[18]), + std::stoi(param_list[19]), + std::stoi(param_list[20]), + param_list[21] == "true" ? true : false, + param_list[22] == "true" ? true : false, + param_list[23] == "true" ? true : false, + sha256}; + }; + + static std::vector loadFromMetaInfoRaw(std::string metainfo_raw) { + std::vector metainfo; + int left_braces = std::count(metainfo_raw.begin(), metainfo_raw.end(), '{'); + int right_braces = std::count(metainfo_raw.begin(), metainfo_raw.end(), '}'); + assert(left_braces == right_braces); + int left_brace_pos = -1; + for (int i = 0; i < metainfo_raw.size(); ++i) { + if (metainfo_raw[i] == '{') { + left_brace_pos = i; + } else if (metainfo_raw[i] == '}') { + metainfo.push_back(TllmGenFmhaKernelMetaInfo::fromString( + metainfo_raw.substr(left_brace_pos + 1, i - left_brace_pos - 1))); + } + } + return metainfo; + }; +}; diff --git a/include/flashinfer/trtllm/fmha/decoder_impl_common.h b/include/flashinfer/trtllm/fmha/decoder_impl_common.h new file mode 100644 index 0000000000000000000000000000000000000000..011a087928ccc08fcffc6955d1767a5d676f3bfb --- /dev/null +++ b/include/flashinfer/trtllm/fmha/decoder_impl_common.h @@ -0,0 +1,247 @@ +/* + * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + * Common utils to be shared between Precompiled and JIT implementation. + */ +#pragma once +// NOTE: we use int32_t sequence lengths as gpt attention plugins use int32_t +// for that. XQA kernels assume all length should use uint32_t. + +#include + +#include "../../utils.cuh" +#include "decoder_params.h" + +// #include "utils.h" + +inline int getMultiProcessorCount() { + int device_id; + int multi_processor_count; + FLASHINFER_CUDA_CALL(cudaGetDevice(&device_id)); + FLASHINFER_CUDA_CALL( + cudaDeviceGetAttribute(&multi_processor_count, cudaDevAttrMultiProcessorCount, device_id)); + return multi_processor_count; +} + +using KVCachePageIndex = int32_t; +using SeqLenDataType = uint32_t; +struct KVCacheListParams { + void const* pool = nullptr; + KVCachePageIndex const* block_indices = + nullptr; // shape: [batchSize][beamWidth][2][maxNbPagesPerSeq]. + SeqLenDataType const* sequence_lengths = + nullptr; // shape: [batchSize][beamWidth] (for compatibility) + // NOTE: max_num_blocks_per_sequence for paged kv cache. + uint32_t capacity = 0; + + KVCacheListParams(void const* _pool, KVCachePageIndex const* _block_indices, + SeqLenDataType const* _sequence_lengths, uint32_t _capacity) + : pool(_pool), + block_indices(_block_indices), + sequence_lengths(_sequence_lengths), + capacity(_capacity) {} + + KVCacheListParams() = default; +}; + +struct XQALaunchParam { + uint32_t num_k_heads; + void* output; + // void const* qkv; + KVCacheListParams kvCacheParams; + uint32_t batch_size; + // float const* kv_scale_quant_orig = nullptr; + float kv_scale_quant_orig = 1.f; + int* cu_seq_lens = nullptr; + uint32_t* semaphores = nullptr; + void* scratch = nullptr; +}; + +struct XQAKernelLoadHashKey { + Data_type data_type; + unsigned int sm; + + bool operator==(XQAKernelLoadHashKey const& other) const { + return data_type == other.data_type && sm == other.sm; + } +}; + +struct XQAKernelLoadHasher { + size_t operator()(XQAKernelLoadHashKey const& s) const { + size_t key = s.data_type; + key <<= 16; + key ^= s.sm; + return key; + } +}; + +struct XQAKernelRuntimeHashKey { + Data_type kv_data_type; + unsigned int head_size; + unsigned int beam_size; + unsigned int num_q_heads_per_kv; + unsigned int m_tilesize; + unsigned int tokens_per_page; + bool paged_kv_cache; + bool multi_query_tokens; + + bool operator==(XQAKernelRuntimeHashKey const& other) const { + return kv_data_type == other.kv_data_type && head_size == other.head_size && + num_q_heads_per_kv == other.num_q_heads_per_kv && beam_size == other.beam_size && + multi_query_tokens == other.multi_query_tokens && m_tilesize == other.m_tilesize && + tokens_per_page == other.tokens_per_page && paged_kv_cache == other.paged_kv_cache; + } +}; +std::ostream& operator<<(std::ostream& os, const XQAKernelRuntimeHashKey& key); + +XQAKernelRuntimeHashKey getRuntimeHashKeyFromXQAParams(XQAParams const& xqaParams); + +void buildXQALaunchParams(XQALaunchParam& launchParams, XQAParams const& params, + KVCacheListParams kv_cache_buffer); + +struct XQAKernelRuntimeHasher { + size_t operator()(XQAKernelRuntimeHashKey const& s) const { + size_t key = s.kv_data_type; + key <<= 16; + key ^= s.head_size; + key <<= 8; + key ^= s.num_q_heads_per_kv; + key <<= 8; + key ^= s.beam_size; + key <<= 6; + key ^= s.m_tilesize; + key <<= 10; + key ^= s.tokens_per_page; + key <<= 1; + key ^= s.paged_kv_cache; + key <<= 1; + key ^= s.multi_query_tokens; + return key; + } +}; + +// XQA kernel can be uniquely identified by (LoadHashKey, RuntimeHashKey). +struct XQAKernelFullHashKey { + XQAKernelLoadHashKey load_key; + XQAKernelRuntimeHashKey runtime_key; + + XQAKernelFullHashKey() = default; + + XQAKernelFullHashKey(XQAKernelLoadHashKey const& load_key, + XQAKernelRuntimeHashKey const& runtime_key) + : load_key(load_key), runtime_key(runtime_key) {} + + XQAKernelFullHashKey(void const* buffer, size_t buffer_size) { + TORCH_CHECK(sizeof(*this) <= buffer_size); + memcpy(this, buffer, sizeof(*this)); + } + + bool operator==(XQAKernelFullHashKey const& other) const { + return load_key == other.load_key && runtime_key == other.runtime_key; + } + + size_t getSerializationSize() const { return sizeof(*this); } + + void serialize(void* buffer, size_t buffer_size) const { + TORCH_CHECK(sizeof(*this) <= buffer_size); + memcpy(buffer, this, sizeof(*this)); + } +}; + +struct XQAKernelFullHasher { + size_t operator()(XQAKernelFullHashKey const& s) const { + return XQAKernelLoadHasher()(s.load_key) ^ XQAKernelRuntimeHasher()(s.runtime_key); + } +}; + +std::uintptr_t constexpr kCudaMemAlign = 128; + +inline int8_t* alignPtr(int8_t* ptr, uintptr_t to) { + uintptr_t addr = (uintptr_t)ptr; + if (addr % to) { + addr += to - addr % to; + } + return (int8_t*)addr; +} + +inline int8_t* nextWorkspacePtrCommon(int8_t* ptr, uintptr_t previousWorkspaceSize, + uintptr_t const alignment) { + uintptr_t addr = (uintptr_t)ptr; + addr += previousWorkspaceSize; + return alignPtr((int8_t*)addr, alignment); +} + +inline int8_t* nextWorkspacePtrWithAlignment(int8_t* ptr, uintptr_t previousWorkspaceSize, + uintptr_t const alignment = kCudaMemAlign) { + return nextWorkspacePtrCommon(ptr, previousWorkspaceSize, alignment); +} + +template +std::optional getGlobalVar(CUmodule hmod, char const* const name, bool required = false) { + T* pVar = nullptr; + size_t size = 0; + auto const error = cuModuleGetGlobal(reinterpret_cast(&pVar), &size, hmod, name); + T ret; + switch (error) { + case CUDA_SUCCESS: + TORCH_CHECK(size == sizeof(T)); + CUDACHECK(cudaMemcpy(&ret, pVar, size, cudaMemcpyDeviceToHost)); + break; + case CUDA_ERROR_NOT_FOUND: + if (!required) { + return std::nullopt; + } + [[fallthrough]]; + default: + cuErrCheck(("Failed to retrieve global variable from cubin.", error)); + } + return std::optional{std::move(ret)}; +} + +inline int computeMultiBlockCount(XQAParams const& xqaParams, int batch_size, + int multiprocessor_count) { + int multi_block_count = 1; + int num_kv_heads = xqaParams.num_kv_heads; + int history_length = xqaParams.timestep; + + int32_t const maxNbSubSeq = kXQA_MAX_NUM_SUB_SEQ; + + multi_block_count = history_length / kMinHistoryTokensPerBlock; + // avoid using too many blocks for one sequence, otherwise the final reduction + // may dominate. + multi_block_count = + std::min(multi_block_count, static_cast(std::round(std::sqrt(multi_block_count * 8.F)))); + multi_block_count = std::max(multi_block_count, 1); + // adjust to kTargetWaveFactor, as already initialized using + // kMinHistoryTokensPerBlock, only need to decrease. + double wave_count = + (double)batch_size * num_kv_heads * multi_block_count / (double)multiprocessor_count; + double adj_factor = wave_count / (double)kTargetWaveFactor; + if (adj_factor > 1.0) { + multi_block_count = floor(multi_block_count / adj_factor); + } + multi_block_count = std::max(multi_block_count, 1); + + // Add limitation due to reserved workspace size. + // When batch_size is large, multi-block is useless anyway. So large workspace + // is not useful and we can set a hard limit for workspace size (computed from + // maxNbSubSeq). + multi_block_count = std::max(std::min(multi_block_count, maxNbSubSeq / batch_size), 1); + + TORCH_CHECK(multi_block_count >= 1, "MultiBlock count should be larger than 1"); + TORCH_CHECK(multi_block_count == 1 || batch_size * multi_block_count <= maxNbSubSeq, + "Insufficient workspace"); + return multi_block_count; +} diff --git a/include/flashinfer/trtllm/fmha/decoder_params.h b/include/flashinfer/trtllm/fmha/decoder_params.h new file mode 100644 index 0000000000000000000000000000000000000000..3d9cc7fc29beb94c7537dbeb0af6bdc3f78fe4ba --- /dev/null +++ b/include/flashinfer/trtllm/fmha/decoder_params.h @@ -0,0 +1,44 @@ +/* + * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once +#include "../common.h" +using XQADataType = Data_type; + +struct XQAParams { + XQADataType data_type = DATA_TYPE_FP16; + XQADataType kv_cache_data_type = DATA_TYPE_FP16; + void* output = nullptr; + void const* qHeads = nullptr; + // float const* kv_scale_quant_orig = nullptr; + float kv_scale_quant_orig = 1.f; + uint32_t* semaphores = nullptr; + void* workspaces = nullptr; + uint32_t batch_size = 0; + int32_t beam_width = 0; + + int32_t num_q_heads = 0; + int32_t num_kv_heads = 0; + int32_t head_size = 0; + int timestep = 0; + + // Paged KV cache parameters. + int generation_input_length; + bool paged_kv_cache = true; // always true + int tokens_per_block; + int max_blocks_per_sequence; + bool multi_block_mode; + bool multi_query_tokens = false; +}; diff --git a/include/flashinfer/trtllm/fmha/fmhaKernels.cuh b/include/flashinfer/trtllm/fmha/fmhaKernels.cuh new file mode 100644 index 0000000000000000000000000000000000000000..c0e34f6b4dcfbbba99f9035dde7e30a148a53fb9 --- /dev/null +++ b/include/flashinfer/trtllm/fmha/fmhaKernels.cuh @@ -0,0 +1,654 @@ +/* + * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +#include +#include +#include +#include +#include + +#include "../../utils.cuh" +#include "../common.h" +#include "cubin/kernelMetaInfo.h" +#include "cuda_runtime_api.h" +#include "fmhaRunnerParams.h" +#include "kernelParams.h" + +#ifdef TLLM_GEN_FMHA_CUBIN_PATH +static const std::string tllm_gen_fmha_cubin_path = std::string(TLLM_GEN_FMHA_CUBIN_PATH); +#else +static_assert(false, "TLLM_GEN_FMHA_CUBIN_PATH macro is not defined when compiling"); +#endif + +#ifdef TLLM_GEN_FMHA_METAINFO_HASH +static const std::string tllm_gen_fmha_metainfo_hash = std::string(TLLM_GEN_FMHA_METAINFO_HASH); +#else +static_assert(false, "TLLM_GEN_FMHA_METAINFO_HASH macro is not defined when compiling"); +#endif + +namespace flashinfer::trtllm_cubin_loader { +std::string getCubin(const std::string& kernelName, const std::string& sha256); +std::string getMetaInfo(const std::string& name, const std::string& sha256, + const std::string& extension); +} // namespace flashinfer::trtllm_cubin_loader +using flashinfer::trtllm_cubin_loader::getCubin; +using flashinfer::trtllm_cubin_loader::getMetaInfo; + +//////////////////////////////////////////////////////////////////////////////////////////////////// +class TllmGenFmhaKernel { + public: + using KernelMeta = TllmGenFmhaKernelMetaInfo; + using RunnerParams = TllmGenFmhaRunnerParams; + using SelectKernelParams = TllmGenSelectKernelParams; + + // Ctor. + TllmGenFmhaKernel(KernelMeta const* pMetaStart, unsigned int nMetaCount, Data_type dtypeQ, + Data_type dtypeKv, Data_type dtypeOut, unsigned int smArch) + : mDtypeQ(dtypeQ), + mDtypeKv(dtypeKv), + mDtypeOut(dtypeOut), + mKernelMeta(pMetaStart), + mKernelMetaCount(nMetaCount), + mSM(smArch) {} + + void loadKernels() { + for (unsigned int i = 0; i < mKernelMetaCount; ++i) { + auto const& kernelMeta = mKernelMeta[i]; + if (kernelMeta.mSM == mSM && kernelMeta.mDataTypeQ == mDtypeQ && + kernelMeta.mDataTypeKv == mDtypeKv && kernelMeta.mDataTypeO == mDtypeOut) { + // Store metadata for later use. + mKernelMetaMap[hashID(kernelMeta)] = i; + } + } + } + + size_t getNumLoadedKernels() const { return mKernelMetaMap.size(); } + + inline uint64_t hashID(int qkvLayout, int maskType, int kernelType, int scheduler, + int multiCtasKvMode, int headDimPerCtaV, int headDimQk, int headDimV, + int tileSizeKv, int numTokensPerPage, int maxNumHeadsQPerKvInCta, + bool reuseSmemKForV, bool uses2CtaMma) const { + TORCH_CHECK((headDimPerCtaV >= 32) && (headDimQk >= 32) && (headDimV >= 32) && + (headDimPerCtaV <= 2048) && (headDimQk <= 2048) && (headDimV <= 2048) && + (numTokensPerPage <= 128), + "Expect (32 <= headDim <= 2048) && (numTokensPerPage <= 128), " + "got headDimPerCtaV=%d, headDimQk=%d, " + "headDimV=%d, numTokensPerPage=%d", + headDimPerCtaV, headDimQk, headDimV, numTokensPerPage); + TORCH_CHECK(maxNumHeadsQPerKvInCta <= 128, "The maxNumHeadsQPerKvInCta <= 128 is required."); + TORCH_CHECK(tileSizeKv == 64 || tileSizeKv == 128, "The tileSizeKv must be 64 or 128."); + // Format of the hash key: + // Bit 0 - 3 : qkvLayout. + // Bit 4 - 7 : maskType. + // Bit 8 - 11: kernelType. + // Bit 12 - 15: tileScheduler. + // Bit 16 - 17: multiCtasKvMode. + // Bit 18 - 24: (headDimPerCtaV >> 5). + // Bit 25 - 31: (headDimQk >> 5). + // Bit 32 - 38: (headDimV >> 5). + // Bit 39 - 40: (tileSizeKv >> 6). + // Bit 41 - 48: numTokensPerPage. + // Bit 49 - 56: maxNumHeadsQPerKvInCta. + // Bit 57 - 57: reuseSmemKForV. + // Bit 58 - 58: uses2CtaMma. + return (static_cast(qkvLayout) << 0) | (static_cast(maskType) << 4) | + (static_cast(kernelType) << 8) | (static_cast(scheduler) << 12) | + (static_cast(multiCtasKvMode) << 16) | + (static_cast(headDimPerCtaV >> 5) << 18) | + (static_cast(headDimQk >> 5) << 25) | + (static_cast(headDimV >> 5) << 32) | + (static_cast(tileSizeKv >> 6) << 39) | + (static_cast(numTokensPerPage) << 41) | + (static_cast(maxNumHeadsQPerKvInCta) << 49) | + (static_cast(reuseSmemKForV) << 57) | + (static_cast(uses2CtaMma) << 58); + } + + uint64_t hashID(KernelMeta const& kernelMeta) const { + return hashID(kernelMeta.mQkvLayout, kernelMeta.mMaskType, kernelMeta.mKernelType, + kernelMeta.mTileScheduler, kernelMeta.mMultiCtasKvMode, + kernelMeta.mHeadDimPerCtaV, kernelMeta.mHeadDimQk, kernelMeta.mHeadDimV, + kernelMeta.mTileSizeKv, kernelMeta.mNumTokensPerPage, + kernelMeta.mMaxNumHeadsQPerKvInCta, kernelMeta.mReuseSmemKForV, + kernelMeta.m2CtaMma); + } + + std::pair checkIfKernelExist(RunnerParams const& params) const { + // The selectKernelParams that might be updated. + SelectKernelParams selectKernelParams{params}; + auto [hashId, info] = hashFromRunnerParams(params, selectKernelParams); + return std::make_pair(mKernelMetaMap.find(hashId) != mKernelMetaMap.end(), info); + } + // start here + void run(RunnerParams const& params) const { + // The selectKernelParams that might be updated. + SelectKernelParams selectKernelParams{params}; + // The iteration index (used to detect a deadlock of selecting new kernels). + int selectKernelIter = 0; + // While loop. + while (true) { + // Any value >= 2 should work here, but we set it larger in case that we + // might have more complicated heuristic in the future. + TORCH_CHECK(selectKernelIter < 8, + "A deadlock is detected when selecting trtllm-gen kernels."); + auto [hashId, info] = hashFromRunnerParams(params, selectKernelParams); + auto const findMetaIter = mKernelMetaMap.find(hashId); + + // Add debug info when kernels are not found. + TORCH_CHECK(findMetaIter != mKernelMetaMap.end(), "Trtllm-gen kernels not found: " + info); + + // auto const& kernelMeta = mKernelMeta[findIter->second.mMetaInfoIndex]; + auto const findFuncIter = mFunctions.find(hashId); + if (findFuncIter == mFunctions.end()) { + // Load the kernel on-demand. + loadKernel(hashId, findMetaIter->second); + } + // Retrieve the loaded kernel. + auto const& kernelInfo = mFunctions.at(hashId); + auto const& kernelMeta = mKernelMeta[kernelInfo.mMetaInfoIndex]; + CUfunction func = kernelInfo.mDeviceFunction; + + // Compute the number of CTAs in X, Y and Z dimension and the cluster size in the X dimension. + auto [maxNumCtasQ, maxNumCtasKv, numCtasX, numCtasY, numCtasZ, clusterDimX] = + computeCtaAndClusterConfig(params, kernelMeta, selectKernelParams); + // Need to select a new kernel if mSelectNewKernel is true. + if (selectKernelParams.mSelectNewKernel) { + selectKernelIter++; + continue; + } + + // Prepare the kernel parameters. + auto kernelParams = + KernelParams::setKernelParams(params, kernelMeta, maxNumCtasQ, maxNumCtasKv); + + // Prepare kernel parameters list for cuLaunchKernelEx. + void* kernelParamsList[] = {&kernelParams}; + CUlaunchConfig launch_config; + launch_config.blockDimX = kernelMeta.mThreadsPerCTA; + launch_config.blockDimY = 1; + launch_config.blockDimZ = 1; + launch_config.gridDimX = numCtasX; + launch_config.gridDimY = numCtasY; + launch_config.gridDimZ = numCtasZ; + launch_config.hStream = params.stream; + launch_config.sharedMemBytes = kernelMeta.mSharedMemBytes; + + // Debug info. + IKL_LOG_DEBUG("TRTLLM-Gen launch info (in TllmGenFmhaKernel %s, %s, %s, %d): kernelName = %s", + toStr(mDtypeQ), toStr(mDtypeKv), toStr(mDtypeOut), mSM, kernelMeta.mFuncName); + IKL_LOG_DEBUG( + "TRTLLM-Gen launch info: maxSeqLenQ = %d, " + "maxSeqLenKv = %d, " + "numHeadsQ = %d, " + "numHeadsKv = %d, batchSize = %d, kernelType = %d", + params.mMaxSeqLenQ, params.mMaxSeqLenKv, params.mNumHeadsQ, params.mNumHeadsKv, + params.mBatchSize, static_cast(params.mKernelType)); + IKL_LOG_DEBUG( + "TRTLLM-Gen launch info: numCtasX = %d, numCtasY = %d, numCtasZ = %d, clusterDimX = %d", + numCtasX, numCtasY, numCtasZ, clusterDimX); + + CUlaunchAttribute launch_attribute[3]; + launch_attribute[0].id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION; + launch_attribute[0].value.clusterDim.x = clusterDimX; + launch_attribute[0].value.clusterDim.y = 1; + launch_attribute[0].value.clusterDim.z = 1; + launch_attribute[1].id = CU_LAUNCH_ATTRIBUTE_CLUSTER_SCHEDULING_POLICY_PREFERENCE; + launch_attribute[1].value.clusterSchedulingPolicyPreference = + clusterDimX > 1 ? CU_CLUSTER_SCHEDULING_POLICY_SPREAD + : CU_CLUSTER_SCHEDULING_POLICY_DEFAULT; + launch_attribute[2].id = CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION; + launch_attribute[2].value.programmaticStreamSerializationAllowed = getEnvEnablePDL(); + + launch_config.attrs = launch_attribute; + launch_config.numAttrs = 3; + // Add setting for non-portable cluster size. + if (clusterDimX > 8) { + cuErrCheck(cuFuncSetAttribute(func, CU_FUNC_ATTRIBUTE_NON_PORTABLE_CLUSTER_SIZE_ALLOWED, + 1 // Enable non-portable cluster sizes + )); + } + + // Force using GmemReduction for the multiCtasKvMode if the CgaSmemReduction needs more than + // one wave (due to the cluster occupancy limit). + // TODO: find a better heuristic of using CgaSmemReduction. + if (isCgaSmemReduction(selectKernelParams.mMultiCtasKvMode)) { + // The maximum number of active clusters that could co-exist. + int maxActiveClusters = 1; + cuErrCheck(cuOccupancyMaxActiveClusters(&maxActiveClusters, func, &launch_config)); + // Use the GmemReduction instead if it needs more than one wave. + if (maxActiveClusters * clusterDimX < (numCtasX * numCtasY * numCtasZ)) { + selectKernelParams.mForceGmemReduction = true; + selectKernelParams.mMultiCtasKvMode = MultiCtasKvMode::GmemReduction; + // continue to select a new kernel. + continue; + } + } + cuErrCheck(cuLaunchKernelEx(&launch_config, func, kernelParamsList, nullptr)); + // Break the while op. + break; + } + } + + private: + // Is it MLA generation kernel ? + inline bool isMlaGenKernel(RunnerParams const& params) const { + return params.mHeadDimQk == 576 && params.mHeadDimV == 512; + } + + // Compute the number of CTAs in X, Y and Z dimension and the cluster size in the X dimension. + using CtaClusterInfo = std::tuple; + + CtaClusterInfo computeCtaAndClusterConfig(RunnerParams const& params, + KernelMeta const& kernelMeta, + SelectKernelParams& selectKernelParams) const { + bool isDsv3MinLatencyMode = params.mBatchSize == 1 && params.mMaxSeqLenQ >= 1 && + params.mMaxSeqLenQ <= 16 && params.mHeadDimQk == 576 && + params.mHeadDimV == 512; + // Do we need to select a new kernel ? + selectKernelParams.mSelectNewKernel = false; + + // The number of Ctas per Q sequence. + int numCtasPerSeqQ = (params.mMaxSeqLenQ + kernelMeta.mStepQ - 1) / kernelMeta.mStepQ; + // Each CTA handles one tokenQ by default for spec-decoding generation kernel, which is used to + // emulate causal masking (like MTP or Eagle3). Note this will be changed later when the + // high-throughput spec-decoding generation kernels are integrated. + if (params.mMaxSeqLenQ > 1 && !isContextKernel(params.mKernelType)) { + numCtasPerSeqQ = params.mMaxSeqLenQ; + } + + // Compute the grid dimension Y. + int numHeadsPerCta = kernelMeta.mGroupsHeadsQ + ? std::min(params.mNumHeadsQPerKv, kernelMeta.mMaxNumHeadsQPerKvInCta) + : 1; + int numCtasForAllHeadsQ = params.mNumHeadsQ / numHeadsPerCta; + TORCH_CHECK(numHeadsPerCta * numCtasForAllHeadsQ == params.mNumHeadsQ, + "The numHeadsQ/numHeadsKv is not supported."); + // Take the number of headDim CTAs. + TORCH_CHECK(kernelMeta.mHeadDimV % selectKernelParams.mHeadDimPerCtaV == 0, + "The headDimPerCtaV is not supported."); + int numCtasPerHeadDim = kernelMeta.mHeadDimV / selectKernelParams.mHeadDimPerCtaV; + // Compute the current numCtasX. + int numCtasX = numCtasPerSeqQ; + // Update the numCtasY. + int numCtasY = numCtasForAllHeadsQ * numCtasPerHeadDim; + // Compute the grid dimension Z. + int numCtasZ = params.mBatchSize; + // The 2CtaMma kernels will use 2 Ctas in the x dimension (only used by MLA generation kernels) + // for heads, so numCtasPerHeadDim and numCtasForAllHeadsQ will be handled by the 2Ctas in the x + // dimension. + if (isMlaGenKernel(params) && selectKernelParams.mUses2CtaMma) { + TORCH_CHECK(numCtasForAllHeadsQ == 2 && numCtasPerHeadDim == 2, + "Internal error: numCtasPerHeadDim should be 2."); + numCtasX *= 2; + numCtasY /= (numCtasForAllHeadsQ * numCtasPerHeadDim); + } + + // First split the seqLenKv into multiple CTAs if the utilization is not full. + // The number of Ctas per KV sequence. + int numCtasPerSeqKv = 1; + // Consider the multiCtasKvMode for better GPU utilization. + if (isMultiCtasKvEnabled(selectKernelParams.mMultiCtasKvMode)) { + // The maximum attention window (the maximum number of tokensKv that will be attended to). + int maxAttentionWindow{params.mMaxSeqLenKv}; + // Some of the tilesKv will be skipped if the sliding window attention or chunked attention is + // used. + if (isSlidingOrChunkedCausalMask(selectKernelParams.mMaskType)) { + if (params.mMaxSeqLenKv > params.mAttentionWindowSize) { + // Consider that the first tileKv might contain tokensKv that is out of the attention + // window. + maxAttentionWindow = + std::min(params.mMaxSeqLenKv, params.mAttentionWindowSize + kernelMeta.mStepKv - 1); + } else { + maxAttentionWindow = std::min(params.mMaxSeqLenKv, params.mChunkedAttentionSize); + } + } + + // The maximum number Ctas per Kv sequence, which makes sure that each CtaKv has work to do. + int const maxNumCtasPerSeqKv = + (maxAttentionWindow + kernelMeta.mStepKv - 1) / kernelMeta.mStepKv; + // Compute numCtasPerSeqKv. + numCtasPerSeqKv = std::min( + maxNumCtasPerSeqKv, + std::max(1, int32_t(params.mMultiProcessorCount / (numCtasX * numCtasY * numCtasZ)))); + // Update the numCtasX. + numCtasX *= numCtasPerSeqKv; + // The current total number of CTAs. + int totalNumCtas = numCtasX * numCtasZ * numCtasY; + // Disable the multiCtasKvMode if there is only one CtaKv. + if (numCtasPerSeqKv <= 1) { + selectKernelParams.mMultiCtasKvMode = MultiCtasKvMode::Disabled; + // Enable the persistent scheduler for better performance. + selectKernelParams.mTileScheduler = TileScheduler::Persistent; + // Need to select a different kernel. + selectKernelParams.mSelectNewKernel = true; + } else if (totalNumCtas < params.mMultiProcessorCount && isMlaGenKernel(params) && + selectKernelParams.mTileSizeKv == 128 && getEnvUseTileSizeKv64ForTrtllmGen()) { + // Use smaller tileSizeKv to fully utilize the SMs. + selectKernelParams.mTileSizeKv = 64; + // Need to select a different kernel. + selectKernelParams.mSelectNewKernel = true; + } + + // Enable the CgaSmemReduction if the numCtasPerSeqKv <= 16 as the maximum cluster dimension + // is 16. Only the swapsMmaAbForGeneration kernel supports the CgaSmemReduction for now. + if (!isDsv3MinLatencyMode && numCtasPerSeqKv > 1 && numCtasPerSeqKv <= 16 && + isSwapsMmaAbForGenerationKernel(selectKernelParams.mKernelType) && + isGmemReduction(selectKernelParams.mMultiCtasKvMode) && + !selectKernelParams.mForceGmemReduction) { + selectKernelParams.mMultiCtasKvMode = MultiCtasKvMode::CgaSmemReduction; + // Need to select a different kernel. + selectKernelParams.mSelectNewKernel = true; + } + + // Add the debug info when multiCtasKvMode is enabled. + if (numCtasPerSeqKv > 1) { + IKL_LOG_DEBUG( + "TRTLLM-Gen launch info: multiCtasKvMode is enabled with tileSizeKv = %d, " + "numCtasPerSeqKv = %d, " + "numCtasPerSeqQ = " + "%d, numCtasY = %d, numCtasZ = %d", + selectKernelParams.mTileSizeKv, numCtasPerSeqKv, numCtasPerSeqQ, numCtasY, numCtasZ); + } + } + + // The cluster size in the X dimension. + int clusterDimX = selectKernelParams.mUses2CtaMma ? 2 : 1; + if (isCgaSmemReduction(selectKernelParams.mMultiCtasKvMode)) { + // Note 2CtaMma and CgaSmemReduction cannot be used together currently. + clusterDimX *= numCtasPerSeqKv; + } + + // Compute the current number of CTAs in total. + int totalNumCtas = numCtasX * numCtasZ * numCtasY; + + // Then split the headDimV into multiple CTAs if there are still unused SMs. + if (isMlaGenKernel(params) && !selectKernelParams.mReuseSmemKForV && + !selectKernelParams.mSelectNewKernel && !selectKernelParams.mUses2CtaMma) { + // Split the headDimV into multiple CTAs if the utilization is not full. + // It doesn't work with reuseSmemKForV currently. + // TODO: find better heuristic of splitting headDimV across multiple CTAs. + + int corrFactor = isDsv3MinLatencyMode ? 1 : 2; + if (selectKernelParams.mHeadDimPerCtaV == 512 && + totalNumCtas * corrFactor <= params.mMultiProcessorCount) { + // Use smaller headDimPerCtaV to fully utilize the SMs. + selectKernelParams.mHeadDimPerCtaV = + totalNumCtas * 2 * corrFactor <= params.mMultiProcessorCount ? 128 : 256; + // Need to select a different kernel. + selectKernelParams.mSelectNewKernel = true; + } + } + + // Return the number of CTAs for X, Y and Z dimension and the cluster size in the X dimension. + return std::make_tuple(numCtasPerSeqQ, numCtasPerSeqKv, numCtasX, numCtasY, numCtasZ, + clusterDimX); + } + + // Determine if we should use the SwapsMmaAbForGeneration kernel for MLA generation. + bool useSwapsMmaAbMlaGenKernel(RunnerParams const& params) const { + // Use the SwapsMmaAbForGeneration kernel for MLA generation when the following conditions are + // met: + // 1. The seqLenPerCtaKv <= 1024 based on the benchmark results (this might be fine-tuned + // later). + // 2. The numCtas (after splitting the heads across multiple CTAs) <= + // params.mMultiProcessorCount. + + // The maximum number Ctas per Kv sequence, which makes sure that each CtaKv has work to do. + // Here we assume the stepKv is 256. + int const maxNumCtasPerSeqKv = flashinfer::ceil_div(params.mMaxSeqLenKv, 256); + ; + // The number of Ctas. + int const numCtas = static_cast(params.mBatchSize * params.mMaxSeqLenQ * + divUp(params.mNumHeadsQPerKv, 16)); + // Compute numCtasPerSeqKv. + int const numCtasPerSeqKv = + std::min(maxNumCtasPerSeqKv, std::max(1, int32_t(params.mMultiProcessorCount / numCtas))); + // Compute the seqLenPerCtaKv. + int const seqLenPerCtaKv = flashinfer::ceil_div(params.mMaxSeqLenKv, numCtasPerSeqKv); + // Whether we should use the SwapsMmaAbForGeneration kernel for MLA generation. + return seqLenPerCtaKv <= 1024 && numCtas <= params.mMultiProcessorCount; + } + + std::pair hashFromRunnerParams( + RunnerParams const& params, SelectKernelParams& selectKernelParams) const { + // The updated kernel type. + FmhaKernelType& kernelType = selectKernelParams.mKernelType; + // Generation kernelType will use either SwapsMmaAbForGeneration or KeepsMmaAbForGeneration. + if (isGenerationKernel(params.mKernelType) && isMlaGenKernel(params)) { + // We use the low-latency kernel (SwapsMmaAbForGeneration with tileSizeQ = 16) when any of the + // following conditions are met: + // 1. The number of headsQPerKv is <= 32. + // 2. The seqLenPerCtaKv <= 1024 based on the benchmark results (this might be fine-tuned + // later) and + // the numCtas (after splitting the heads across multiple CTAs) <= + // params.mMultiProcessorCount. + + // Check the conditions. + if (params.mNumHeadsQPerKv <= 32 || useSwapsMmaAbMlaGenKernel(params)) { + kernelType = FmhaKernelType::SwapsMmaAbForGeneration; + } else { + // Otherwise, we use the high-throughput kernel. + kernelType = FmhaKernelType::KeepsMmaAbForGeneration; + // The 2CTA keepsMmaAbForGeneration kernel is used when the numHeadsQPerKv is 128. + if (params.mNumHeadsQPerKv == 128) { + selectKernelParams.mUses2CtaMma = true; + // Each Cta only handles 256 headDimV. + selectKernelParams.mHeadDimPerCtaV = 256; + } + } + } else if (isGenerationKernel(params.mKernelType)) { + kernelType = (params.mNumHeadsQPerKv <= 16 && params.mHeadDimQk != 32) + ? FmhaKernelType::SwapsMmaAbForGeneration + : FmhaKernelType::KeepsMmaAbForGeneration; + } + + // The maximum number of headsQPerKv that the kernel can support in one Cta. + int maxNumHeadsQPerKvInCta = 1; + if (isSwapsMmaAbForGenerationKernel(kernelType)) { + // Set the corresponding maxNumHeadsQPerKvInCta (tileSizeQ) for low-latency generation + // kernels. + maxNumHeadsQPerKvInCta = (params.mNumHeadsQPerKv <= 8) ? 8 : 16; + TORCH_CHECK((maxNumHeadsQPerKvInCta == 8 || maxNumHeadsQPerKvInCta == 16) && + (params.mNumHeadsQPerKv < maxNumHeadsQPerKvInCta || + params.mNumHeadsQPerKv % maxNumHeadsQPerKvInCta == 0), + "Not supported"); + } else if (isKeepsMmaAbForGenerationKernel(kernelType)) { + // Use the maxNumHeadsQPerKvInCta (tileSizeQ) = 64 for MLA high-throughput generation kernels. + maxNumHeadsQPerKvInCta = isMlaGenKernel(params) ? 64 : 32; + TORCH_CHECK((params.mNumHeadsQPerKv < maxNumHeadsQPerKvInCta || + params.mNumHeadsQPerKv % maxNumHeadsQPerKvInCta == 0), + "Not supported"); + } else if (isContextKernel(kernelType)) { + TORCH_CHECK(maxNumHeadsQPerKvInCta == 1, "Not supported"); + } + + // The mask type. + selectKernelParams.mMaskType = params.mMaskType; + // Enable sliding window or chunked causal if the max kv sequence length exceeds attention + // window size or chunked attention size. This is supported by causal-mask context kernels and + // generation-phase kernels. + if ((selectKernelParams.mMaskType == TrtllmGenAttentionMaskType::Causal || + !isContextKernel(params.mKernelType)) && + (params.mMaxSeqLenKv > params.mAttentionWindowSize || + params.mChunkedAttentionSize != INT_MAX)) { + TORCH_CHECK(params.mMaxSeqLenKv <= params.mAttentionWindowSize || + params.mMaxSeqLenKv <= params.mChunkedAttentionSize, + "Sliding window attention and chunked attention should not be used together"); + selectKernelParams.mMaskType = TrtllmGenAttentionMaskType::SlidingOrChunkedCausal; + } + // NumTokensPerPage is set to 0 when not selecting pagedKv-layout kernels. + int numTokensPerPage = (!isPagedKv(params.mQkvLayout)) ? 0 : params.mNumTokensPerPage; + + // Debug info. + std::string info = + "qkvLayout=" + std::to_string(static_cast(params.mQkvLayout)) + + ", maskType=" + std::to_string(static_cast(selectKernelParams.mMaskType)) + + ", kernelType=" + std::to_string(static_cast(kernelType)) + + ", tileScheduler=" + std::to_string(static_cast(selectKernelParams.mTileScheduler)) + + ", multiCtasKvMode=" + + std::to_string(static_cast(selectKernelParams.mMultiCtasKvMode)) + + ", headDimPerCtaV=" + std::to_string(selectKernelParams.mHeadDimPerCtaV) + + ", headDimQk=" + std::to_string(params.mHeadDimQk) + + ", headDimV=" + std::to_string(params.mHeadDimV) + + ", tileSizeKv=" + std::to_string(selectKernelParams.mTileSizeKv) + + ", numTokensPerPage=" + std::to_string(numTokensPerPage) + + ", maxNumHeadsQPerKvInCta=" + std::to_string(maxNumHeadsQPerKvInCta) + + ", reuseSmemKForV=" + std::to_string(selectKernelParams.mReuseSmemKForV) + + ", uses2CtaMma=" + std::to_string(selectKernelParams.mUses2CtaMma); + IKL_LOG_DEBUG( + "Searching for kernel traits (%d available) in TllmGenFmhaKernel(%s, %s, %s, %d) %s", + getNumLoadedKernels(), toStr(mDtypeQ), toStr(mDtypeKv), toStr(mDtypeOut), mSM, + info.c_str()); + + return std::make_pair( + hashID(static_cast(params.mQkvLayout), static_cast(selectKernelParams.mMaskType), + static_cast(kernelType), static_cast(selectKernelParams.mTileScheduler), + static_cast(selectKernelParams.mMultiCtasKvMode), + selectKernelParams.mHeadDimPerCtaV, params.mHeadDimQk, params.mHeadDimV, + selectKernelParams.mTileSizeKv, numTokensPerPage, maxNumHeadsQPerKvInCta, + selectKernelParams.mReuseSmemKForV, selectKernelParams.mUses2CtaMma), + info); + } + + // Load a single kernel (called by `run()` when needed). + void loadKernel(uint64_t hashId, unsigned int metaIndex) const { + auto const& kernelMeta = mKernelMeta[metaIndex]; + CUmodule hmod{0}; + std::string kernelName(kernelMeta.mFuncName); + + // Check if the module is already loaded. + auto findModuleIter = mModules.find(kernelMeta.mFuncName); + auto capitalizeFirst = [](std::string str) { + if (!str.empty()) { + str[0] = std::toupper(str[0]); + } + return str; + }; + if (findModuleIter == mModules.end()) { + // Load the module. + std::string cubin_path = tllm_gen_fmha_cubin_path + kernelMeta.mFuncName; + std::string cubin = getCubin(cubin_path, kernelMeta.sha256); + if (cubin.empty()) { + throw std::runtime_error("Failed to load cubin for " + kernelName); + } + cuErrCheck(cuModuleLoadData(&hmod, cubin.data())); + mModules[kernelName] = hmod; + } else { + hmod = findModuleIter->second; + } + + // Load the function. + KernelInfo funcInfo; + funcInfo.mMetaInfoIndex = metaIndex; + cuErrCheck(cuModuleGetFunction(&funcInfo.mDeviceFunction, hmod, kernelMeta.mFuncName)); + + if (kernelMeta.mSharedMemBytes >= 48 * 1024) { + cuErrCheck(cuFuncSetAttribute(funcInfo.mDeviceFunction, + CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, + kernelMeta.mSharedMemBytes)); + } + + // Cache the loaded function. + mFunctions[hashId] = funcInfo; + } + + Data_type mDtypeQ, mDtypeKv, mDtypeOut; + KernelMeta const* mKernelMeta; + unsigned int mKernelMetaCount; + unsigned int mSM; + mutable std::unordered_map mModules; + + mutable std::unordered_map mKernelMetaMap; + + struct KernelInfo { + unsigned int mMetaInfoIndex; + CUfunction mDeviceFunction; + }; + + mutable std::unordered_map mFunctions; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +class TllmFmhaKernelFactory { + public: + using KernelType = TllmGenFmhaKernel; + + KernelType const* getKernels(Data_type dtypeQ, Data_type dtypeKv, Data_type dtypeOut, + unsigned int sm) { + static std::mutex s_mutex; + std::lock_guard lg(s_mutex); + + if (!metainfo_loaded) { + std::string metainfo_raw = getMetaInfo(tllm_gen_fmha_cubin_path + "flashInferMetaInfo", + tllm_gen_fmha_metainfo_hash, ".h"); + metainfo = KernelType::KernelMeta::loadFromMetaInfoRaw(metainfo_raw); + metainfo_loaded = true; + } + + auto const id = hashID(dtypeQ, dtypeKv, dtypeOut, sm); + auto const findIter = mKernels.find(id); + if (findIter == mKernels.end()) { + KernelType* newKernel = + new KernelType{metainfo.data(), metainfo.size(), dtypeQ, dtypeKv, dtypeOut, sm}; + newKernel->loadKernels(); + mKernels.insert(std::make_pair(id, std::unique_ptr(newKernel))); + IKL_LOG_DEBUG( + "Loading new kernel for dtypeQ=%s, dtypeKv=%s, dtypeOut=%s, sm=%d with %d loaded kernels", + toStr(dtypeQ), toStr(dtypeKv), toStr(dtypeOut), sm, newKernel->getNumLoadedKernels()); + return newKernel; + } + return findIter->second.get(); + } + + static TllmFmhaKernelFactory& Get() { + int deviceId; + cudaGetDevice(&deviceId); + static std::unique_ptr sFactory[32] = {nullptr}; + if (sFactory[deviceId] == nullptr) { + TORCH_CHECK(deviceId < 32, "Invalid deviceId %d (max is 32 devices)", deviceId); + sFactory[deviceId] = std::make_unique(TllmFmhaKernelFactory()); + } + + return *(sFactory[deviceId]); + } + + private: + TllmFmhaKernelFactory() = default; + + inline uint64_t hashID(Data_type dtypeQ, Data_type dtypeKv, Data_type dtypeOut, + unsigned int sm) const { + return static_cast(sm) | static_cast(dtypeQ) << 16 | + static_cast(dtypeKv) << 20 | static_cast(dtypeOut) << 24; + } + + std::unordered_map> mKernels; + std::vector metainfo; + bool metainfo_loaded = false; +}; + +inline TllmGenFmhaKernel const* getTllmFmhaKernels(Data_type dtypeQ, Data_type dtypeKv, + Data_type dtypeOut, unsigned int sm) { +#ifndef EXCLUDE_SM_100 + return TllmFmhaKernelFactory::Get().getKernels(dtypeQ, dtypeKv, dtypeOut, sm); +#else + return nullptr; +#endif // EXCLUDE_SM_100 +} diff --git a/include/flashinfer/trtllm/fmha/fmhaRunner.cuh b/include/flashinfer/trtllm/fmha/fmhaRunner.cuh new file mode 100644 index 0000000000000000000000000000000000000000..f3488cd5ab5799f790437223dad48c8977fc68fc --- /dev/null +++ b/include/flashinfer/trtllm/fmha/fmhaRunner.cuh @@ -0,0 +1,48 @@ +/* + * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +#include "fmhaKernels.cuh" +#include "fmhaRunnerParams.h" + +class TllmGenFmhaRunner { + public: + // Constructor. + explicit TllmGenFmhaRunner(Data_type dtypeQ, Data_type dtypeKv, Data_type dtypeOut); + + TllmGenFmhaRunner() = default; + + // Check if fmha is supported. + bool isSupported(TllmGenFmhaRunnerParams const& runnerParams) const; + + // Check if fmha is supported with additional info. + std::pair isSupportedWithInfo( + TllmGenFmhaRunnerParams const& runnerParams) const; + + // Run the fmha kernel. + void run(TllmGenFmhaRunnerParams const&); + + private: + // The input/output datatype. + Data_type mDtypeQ, mDtypeKv, mDtypeOut; + // The SM version. + int mSM; + // The class that stores all the kernels. + TllmGenFmhaKernel const* mKernel; +}; diff --git a/include/flashinfer/trtllm/fmha/fmhaRunnerParams.h b/include/flashinfer/trtllm/fmha/fmhaRunnerParams.h new file mode 100644 index 0000000000000000000000000000000000000000..738b2a3f74efeb7473ec59cb6e0c7a723034621b --- /dev/null +++ b/include/flashinfer/trtllm/fmha/fmhaRunnerParams.h @@ -0,0 +1,352 @@ +/* + * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +#include +#include +#include + +#include "flashinfer/exception.h" + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// The attention mask types. +enum class TrtllmGenAttentionMaskType { + // Dense mask. + Dense = 0, + // Causal mask. + Causal, + // Sliding window or chunked causal mask. + SlidingOrChunkedCausal, + // Custom mask. + Custom +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// Helper functions to check the mask type. + +#define ATTENTION_MASK_TYPE_FUNCTION(MaskType) \ + inline bool is##MaskType##Mask(TrtllmGenAttentionMaskType maskType) { \ + return (maskType == TrtllmGenAttentionMaskType::MaskType); \ + } + +ATTENTION_MASK_TYPE_FUNCTION(Dense) +ATTENTION_MASK_TYPE_FUNCTION(Causal) +ATTENTION_MASK_TYPE_FUNCTION(SlidingOrChunkedCausal) +ATTENTION_MASK_TYPE_FUNCTION(Custom) + +#undef ATTENTION_MASK_TYPE_FUNCTION + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +enum class FmhaKernelType { + // The context-phase kernels. + Context = 0, + // Choose the best generation kernel based on the heuristic: + // use SwapsMmaAbForGeneration kernels when numHeadsQPerKv <= 16, otherwise + // KeepsMmaAbForGeneration. + Generation = 1, + // Swap tensor A and tensor B of Mma, which only supports numHeadsQPerKv <= 16. + SwapsMmaAbForGeneration, + // Keep tensor A and tensor B of Mma. + KeepsMmaAbForGeneration, + // Speculative decoding (Medusa and Eagle) generation-phase attention kernels, where seqLenQ > 1. + SpecDecodingGeneration +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// Helper functions to check the fmha kernel type. + +#define FMHA_KERNEL_TYPE_FUNCTION(KernelType) \ + inline bool is##KernelType##Kernel(FmhaKernelType kernelType) { \ + return (kernelType == FmhaKernelType::KernelType); \ + } + +FMHA_KERNEL_TYPE_FUNCTION(Context) +FMHA_KERNEL_TYPE_FUNCTION(Generation) +FMHA_KERNEL_TYPE_FUNCTION(SwapsMmaAbForGeneration) +FMHA_KERNEL_TYPE_FUNCTION(KeepsMmaAbForGeneration) +FMHA_KERNEL_TYPE_FUNCTION(SpecDecodingGeneration) + +#undef QKV_LAYOUT_FUNCTION + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// Note that (batchSize, seqLen) dimensions will be packed as sumOfSeqLens without paddings for +// variable sequence lengths. +enum class QkvLayout { + // SeparateQkv: separate Q, K and V buffers. + // Each has the shape: [batchSize, seqLen, numHeads, headDim]. + SeparateQkv = 0, + // PackedQkv: single buffer for Q, K and V. + // Shape: [batchSize, seqLen, numHeadsQ + 2*numHeadsKv, headDim]. + PackedQkv, + // Paged buffer for K and V. Its shape is [batchSize, 2, maxNumPagesPerSeq]. The 2 corresponds to + // K + // and V. That buffer stores the logical page index of the paged-KV memory pool. Each "page" of + // that + // pool is a contiguous buffer of shape [numHeadsKv, pageSize, headDim]. + PagedKv, + // ContiguousKv: + // Contiguous buffer for Q with shape [batchSize, seqLen, numHeads, headDim]. + // Contiguous buffer for Kv with shape [batchSize, seqLen, 2 * numHeads, headDim]. + ContiguousKv, +}; + +// Helper functions to check the QkvLayout type. + +#define QKV_LAYOUT_FUNCTION(LayoutType) \ + inline bool is##LayoutType(QkvLayout qkvLayout) { return (qkvLayout == QkvLayout::LayoutType); } + +QKV_LAYOUT_FUNCTION(SeparateQkv) +QKV_LAYOUT_FUNCTION(PackedQkv) +QKV_LAYOUT_FUNCTION(PagedKv) +QKV_LAYOUT_FUNCTION(ContiguousKv) + +#undef QKV_LAYOUT_FUNCTION + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +enum class TileScheduler { + // Static scheduler (Non-persistent). + Static = 0, + // Persistent scheduler. + Persistent +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// +enum class MultiCtasKvMode { + // No multiCtasKvMode. + Disabled = 0, + // Do the reduction through the global memory and atomic counters. + GmemReduction, + // Do the reduction through the CGA remote shared memory. + CgaSmemReduction +}; + +// Helper function to check if the multiCtasKv is enabled. +inline bool isMultiCtasKvEnabled(MultiCtasKvMode multiCtasKvMode) { + return multiCtasKvMode != MultiCtasKvMode::Disabled; +} + +// Helper function to check the multiCtasKvMode type. + +#define MULTI_CTAS_KV_MODE_FUNCTION(Type) \ + inline bool is##Type(MultiCtasKvMode multiCtasKvMode) { \ + return (multiCtasKvMode == MultiCtasKvMode::Type); \ + } + +MULTI_CTAS_KV_MODE_FUNCTION(Disabled) +MULTI_CTAS_KV_MODE_FUNCTION(GmemReduction) +MULTI_CTAS_KV_MODE_FUNCTION(CgaSmemReduction) + +#undef MULTI_CTAS_KV_MODE_FUNCTION + +//////////////////////////////////////////////////////////////////////////////////////////////////// +struct TllmGenFmhaRunnerParams { + // Input layout. + QkvLayout mQkvLayout; + // Attention mask type. + TrtllmGenAttentionMaskType mMaskType; + // The kernel type. + FmhaKernelType mKernelType; + // The tile scheduler. + TileScheduler mTileScheduler; + // The multiCtasKvMode (i.e. multiBlockMode). + bool mMultiCtasKvMode; + + // Input QKV buffers. + void const* qPtr; + void const* kPtr; + void const* vPtr; + // Packed KV buffer + void const* kvPtr; + // Packed QKV buffer + void const* qkvPtr; + // The scaling factor pointer of K. + void const* kSfBasePtr; + // The scaling factor pointer of V. + void const* vSfBasePtr; + // The custom mask ptr. + uint32_t const* customMaskPtr; + // The packed custom mask's offsets of each sequence. + int64_t const* customMaskOffsetsPtr; + // The first sparseMask offsets in the Kv sequence dimension. + int32_t const* firstSparseMaskOffsetsKvPtr; + // The counter for the multiCtasKv mode. + int32_t* multiCtasKvCounterPtr; + // The sequence length buffer for K/V. + int const* seqLensKvPtr; + // The cumulative sequence length buffer for Q and K/V + int const* cumSeqLensQPtr; + int const* cumSeqLensKvPtr; + // The kv page idx + int const* kvPageIdxPtr; + bool useGmemScale; + // The device output scale for FP8 quantization. + float const* outputScalePtr; + float outputScale; + // The device scaling factor for softmax (multiplied by log2 to use faster exp2) + float const* scaleSoftmaxLog2Ptr; + float scaleSoftmaxLog2; + // The device scale for KV scaling factor. + float const* kvSfScalePtr; + // The device scale for O scaling factor. + float const* oSfScalePtr; + // The scratch space for each CtaKv when the multiCtasKv mode is enabled. + // PartialO, partialMax and partialSum will be stored to the scratch space. + void* multiCtasKvScratchPtr; + // The softmax stats buffer. + // The softmax max/sum values will be stored to the buffer if it is not nullptr. + float2* softmaxStatsPtr; + // Attention sink + float const* ptrAttentionSinks{nullptr}; + // The output buffer. + void* oPtr; + // The output scaling factor buffer. + void* oSfPtr; + + // KV-Cache strides + // The stride between different keys/vals. + int kvStrideKeysValues; + // The stride between different heads. + int kvStrideHeads; + // The stride between different batches. + int kvStrideBatch; + + // Head dimension for Q and K. + int mHeadDimQk; + // Head dimension for V. + int mHeadDimV; + // Number of heads for Q and K/V. + int mNumHeadsQ, mNumHeadsKv, mNumHeadsQPerKv; + // The batch size. + int mBatchSize; + // The max sequence length in the contiguous Kv cache. + int mMaxSeqLenCacheKv; + // The max q sequence length. + int mMaxSeqLenQ; + // The max kv sequence length. + int mMaxSeqLenKv; + // The attention window size for sliding window attention (sliding-window-attention is enabled + // when seqLenKv > mAttentionWindowSize). + int mAttentionWindowSize; + // The chunked attention size (chunked-context is enabled when seqLenKv > mChunkedAttentionSize). + int mChunkedAttentionSize; + // The sum of sequence lengths for Q and K/V. (Only used when mSupportsVarSeqLens = true) + int mSumOfSeqLensQ; + int mSumOfSeqLensKv; + // The maximum number of pages per sequence in the paged-kv buffer. + int mMaxNumPagesPerSeqKv; + // The number of tokens per pageKv. + int mNumTokensPerPage; + // The number of pages in memory pool. + int mNumPagesInMemPool; + // The number of multiProcessor for the GPU. + int mMultiProcessorCount; + // Scaling factor for Q. + float mScaleQ; + // Scaling factor for output. + float mScaleOutput; + // The start token index in SF tensor. Used for FP4 SF offset calculation in generation phase + // kernel when inflight batching is enabled. + int mSfStartTokenIdx; + // The SF scale for Kv. + float mScaleSfKv; + // The SF scale for output. + float mScaleSfO; + // The cuda stream. + cudaStream_t stream; + + // set the attention mask type + TllmGenFmhaRunnerParams& setAttentionMaskType(std::int8_t maskType) { + // maskType is the enum of tensorrt_llm::kernels::ContextAttentionMaskType + // convert ContextAttentionMaskType to TrtllmGenAttentionMaskType + switch (maskType) { + case 0: // tensorrt_llm::kernels::ContextAttentionMaskType::PADDING + mMaskType = TrtllmGenAttentionMaskType::Dense; + break; + case 1: // tensorrt_llm::kernels::ContextAttentionMaskType::CAUSAL + mMaskType = TrtllmGenAttentionMaskType::Causal; + break; + case 2: // tensorrt_llm::kernels::ContextAttentionMaskType::SLIDING_OR_CHUNKED_CAUSAL + mMaskType = TrtllmGenAttentionMaskType::SlidingOrChunkedCausal; + break; + case 3: // tensorrt_llm::kernels::ContextAttentionMaskType::CUSTOM_MASK + mMaskType = TrtllmGenAttentionMaskType::Custom; + break; + default: + FLASHINFER_ERROR("Invalid attention mask type"); + } + return *this; + } + + TllmGenFmhaRunnerParams() { + // NOTE(Zihao): all fields are POD types, so we can use memset to initialize them to zero + static_assert(std::is_standard_layout::value, + "TllmGenFmhaRunnerParams must be a POD type (standard layout) for memset to be " + "safe."); + memset(this, 0, sizeof(TllmGenFmhaRunnerParams)); + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// Parameters that might be updated when selecting kernels. + +struct TllmGenSelectKernelParams { + // The FMHA kernel type. + FmhaKernelType mKernelType; + // The headDimV per CTA, which is only used by MLA generation kernels currently. + int mHeadDimPerCtaV; + // The multiCtasKvMode. + MultiCtasKvMode mMultiCtasKvMode; + // Force using GmemRedution for the multiCtasKvMode. + bool mForceGmemReduction; + // The mask type. + TrtllmGenAttentionMaskType mMaskType; + // Reuse smemK for V or not (only work with MLA generation kernels). + bool mReuseSmemKForV; + // Do we need to select a new kernel as the parameters have been updated. + bool mSelectNewKernel; + // The tile scheduler. + TileScheduler mTileScheduler; + // The tile size for Kv. + int mTileSizeKv; + // Use 2 CTA MMA or not. + bool mUses2CtaMma; + + // The constructor. + TllmGenSelectKernelParams(TllmGenFmhaRunnerParams params) + : mKernelType(params.mKernelType), + mHeadDimPerCtaV(params.mHeadDimV) + // Note the CgaSmemReduction will be enabled based on the heuristic. + , + mMultiCtasKvMode(params.mMultiCtasKvMode ? MultiCtasKvMode::GmemReduction + : MultiCtasKvMode::Disabled), + mForceGmemReduction(false), + mMaskType(params.mMaskType), + mReuseSmemKForV(false), + mSelectNewKernel(false), + mTileScheduler(params.mTileScheduler), + mTileSizeKv(128), + mUses2CtaMma(false) {}; +}; diff --git a/include/flashinfer/trtllm/fmha/gen_kernel_launcher.cuh b/include/flashinfer/trtllm/fmha/gen_kernel_launcher.cuh new file mode 100644 index 0000000000000000000000000000000000000000..0bf4d891f84dd90714f6300364e3cbd997829e03 --- /dev/null +++ b/include/flashinfer/trtllm/fmha/gen_kernel_launcher.cuh @@ -0,0 +1,21 @@ +/* + * Copyright (c) 2020-2024, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once +#include +#include + +#include "decoder_impl_common.h" +#include "decoder_params.h" diff --git a/include/flashinfer/trtllm/fmha/kernelParams.h b/include/flashinfer/trtllm/fmha/kernelParams.h new file mode 100644 index 0000000000000000000000000000000000000000..b92e24db0066db3e7b7393370402e69e3a2f1535 --- /dev/null +++ b/include/flashinfer/trtllm/fmha/kernelParams.h @@ -0,0 +1,698 @@ +/*************************************************************************************************** + * Copyright (c) 2011-2024, NVIDIA CORPORATION. All rights reserved. + * + * Redistribution and use in source and binary forms, with or without modification, are not permit- + * ted. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR + * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND + * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, + * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; + * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, + * STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#ifdef _WIN32 +#define _USE_MATH_DEFINES +#include +#endif + +#include + +#include +#include +#include + +#include "../common.h" +#include "fmhaRunnerParams.h" + +//////////////////////////////////////////////////////////////////////////////////////////////////// +using Dtype = Data_type; + +struct KernelParams { + // TMA descriptor for Q. + CUtensorMap tmaQ_; + // TMA descriptor for K. + CUtensorMap tmaK_; + // TMA descriptor for V. + CUtensorMap tmaV_; + // The descriptor for O. + CUtensorMap tmaO_; + + // For FP4 KV cache, additional scaling factors are needed. + // TMA descriptor for K scaling factor. + CUtensorMap tmaKSf_; + // TMA descriptor for V scaling factor. + CUtensorMap tmaVSf_; + + // grid dimensions, these might differ from actual grid the kernel is launched with + // for persistent kernels on Hopper GPUs. + int32_t logicalGridDimX, logicalGridDimY, logicalGridDimZ; + + // The output pointer (used by STG for last tile). + void* ptrO; + // The output SF pointer (used for FP4 output). + void* ptrSfO; + // The attention sinks pointer (additional value per head in the denominator of the softmax). + float const* ptrAttentionSinks; + + // The cumulative sequence lengths for Q. + int32_t const* ptrCumSeqLensQ; + // The cumulative sequence lengths for K/V. + int32_t const* ptrCumSeqLensKv; + // The packed custom mask. + uint32_t const* ptrCustomMask; + // The packed custom mask's offsets of each sequence. + int64_t const* ptrCustomMaskOffsets; + // The debug output matrix O + float* ptrDebugO; + // The first sparseMask offsets in the Kv sequence dimension. + int32_t const* ptrFirstSparseMaskOffsetsKv; + // The counter for the multiCtasKv mode. + int32_t* ptrMultiCtasKvCounter; + // The device output scale for FP8 quantization. Only needed by trt-llm fp8 kernels as the sca- + // les have to be on the device currently. + float const* ptrOutputScale; + // The page indexes of the paged-kv buffer with shape of [batchSize, 2, maxNumPagesPerSeq]. + int32_t const* ptrPageIdxKv; + // The partial matrix O for each CtaKv when the multiCtasKv mode is enabled. + void* ptrPartialO; + // The partial softmax stats (max/sum)for each CtaKv when the multiCtasKv mode is enabled. + float2* ptrPartialStats; + // The scaling factors for K. + float const* ptrSageAttnSfsK; + // The scaling factors for P. + float const* ptrSageAttnSfsP; + // The scaling factors for Q. + float const* ptrSageAttnSfsQ; + // The scaling factors for V. + float const* ptrSageAttnSfsV; + // The device scaling factor for softmax (multiplied by log2 to use faster exp2). Only needed by + // trt-llm fp8 kernels as the scales have to be on the device currently. + float const* ptrScaleSoftmaxLog2; + // The SF scale for Kv on device. Only needed by trt-llm kernels as the scales have to be on the + // device currently. + float const* ptrScaleSfKv; + // The SF scale for O on device. Only needed by trt-llm kernels as the scales have to be on the + // device currently. + float const* ptrScaleSfO; + // The sequence lengths for K/V. Required by pagedKv kernels to avoid unnecessary computation + // based on (ptrCumSeqLensKv[batchIdx + 1] - ptrCumSeqLensKv[batchIdx]). + int32_t const* ptrSeqLensKv; + // The softmax stats buffer. + float2* ptrSoftmaxStats; + + // The attention window size for sliding window attention. + int32_t mAttentionWindowSize; + // The batch size + int32_t mBatchSize; + // The chunked attention size in log2. + int32_t mChunkedAttentionSizeLog2; + // The log of the Sage Attention block size for K. + int32_t mLogNumEltsPerSageAttnBlkK; + // The log of the Sage Attention block size for P. + int32_t mLogNumEltsPerSageAttnBlkP; + // The log of the Sage Attention block size for Q. + int32_t mLogNumEltsPerSageAttnBlkQ; + // The log of the Sage Attention block size for V. + int32_t mLogNumEltsPerSageAttnBlkV; + // The sequence lengths for Q and K/V. + int32_t mMaxSeqLenQ, mMaxSeqLenKv; + // The maximum number of CTAs for Q. + int32_t mMaxNumCtasQ; + // The maximum number of CTAs for K/V. + int32_t mMaxNumCtasKv; + // The maximum number of pages per sequence for paged-kv buffer. + int32_t mMaxNumPagesPerSeqKv; + // The number of heads for K/V. + int32_t mNumHeadsKv; + // The number of heads for Q. + int32_t mNumHeadsQ; + // The number of Q heads per K/V head (i.e. mNumHeadsQ / mNumHeadsKv). + int32_t mNumHeadsQPerKv; + // The hidden size of O. + int64_t mNumHiddenEltsO; + // The number of MTP tokens per sequence. Assume that all requests have the same numMtpTokens + // without paddings. + int32_t mNumMtpTokens; + // The total number of pages in the paged-kv memory pool. + int32_t mNumPagesInMemPool; + // The output scale for FP8 quantization. + float mOutputScale; + // The scaling factor for softmax (multiplied by log2 to use faster exp2). + float mScaleSoftmaxLog2; + // The SF scale for Kv. + float mScaleSfKv; + // The SF scale for O. + float mScaleSfO; + // The start token index in SF tensor. Used for FP4 SF offset calculation in generation phase + // kernel when inflight batching is enabled in TRT-LLM. + int32_t mStartTokenIdxSfO; + // The sum of sequence lengths for Q and K/V. + int32_t mSumOfSeqLensQ, mSumOfSeqLensKv; + + // Create the TMA shape/stride for Q. + template + static auto makeTmaShapeStrideQ(FmhaOptions const& options, bool groupsHeadsQ, int32_t tileSizeQ, + int32_t numEltsInClampedHeadDimQ) { + // + // The Q has shape of [numTokens * numHeadsQPerKv, numHeadsKv * 1, headDim] + // when grouping headsQ, otherwise it would be [numTokens, numHeadsQPerKv * numHeadsKv, + // headDim]. + + // The number of grouped heads for the A matrix of MMA. + int32_t numGroupedHeads{1}; + if (groupsHeadsQ) { + numGroupedHeads = std::min(tileSizeQ, options.mNumHeadsQPerKv); + } + + // The number of heads. + int32_t numHeads{options.mNumHeadsQ}; + if (groupsHeadsQ) { + numHeads /= numGroupedHeads; + } + // Make sure the math works. + TORCH_CHECK(numHeads * numGroupedHeads == options.mNumHeadsQ, "internal error"); + + // The number of tokens. + int32_t numTokens{options.mSumOfSeqLensQ}; + + // This maps to flattened TMA shape for Q: (headDim, numTokens, numHeads). + auto shape = std::vector{ + static_cast(options.mHeadDimQk), static_cast(numGroupedHeads), + static_cast(numHeads), static_cast(numTokens)}; + + // The hidden dimension when the tensor contains only Q (i.e. not QKV packed). + int32_t const hiddenDimQ{options.mNumHeadsQ * options.mHeadDimQk}; + + // The hidden dimension when the Q, K and V tensors are packed. + int32_t hiddenDimQkv{hiddenDimQ}; + if (isPackedQkv(options.mQkvLayout)) { + TORCH_CHECK(!groupsHeadsQ, "internal error"); + hiddenDimQkv += options.mNumHeadsKv * (options.mHeadDimQk + options.mHeadDimV); + } + + // The stride between tokens. + int32_t strideTokens{hiddenDimQkv}; + + // The stride between heads. + int32_t strideHeads{groupsHeadsQ ? numGroupedHeads * options.mHeadDimQk : options.mHeadDimQk}; + + // The stride between grouped heads. + int32_t strideGroupedHeads{options.mHeadDimQk}; + + // Assemble the stride (1, strideTokens, strideHeads). + // Swap the first two dimension as mentioned before. + auto stride = std::vector{1, static_cast(strideGroupedHeads), + static_cast(strideHeads), + static_cast(strideTokens)}; + + // The tile shape for TMA. + auto tileShapes = std::vector{static_cast(numEltsInClampedHeadDimQ), 1, 1, + static_cast(tileSizeQ)}; + if (groupsHeadsQ) { + if (isSpecDecodingGenerationKernel(options.mKernelType)) { + TORCH_CHECK((tileSizeQ % numGroupedHeads == 0), "internal error"); + tileShapes = std::vector{static_cast(numEltsInClampedHeadDimQ), + static_cast(numGroupedHeads), 1, + static_cast(tileSizeQ / numGroupedHeads)}; + } else { + tileShapes = std::vector{static_cast(numEltsInClampedHeadDimQ), + static_cast(tileSizeQ), 1, 1}; + } + } + + return std::make_tuple(shape, stride, tileShapes); + } + + // Create the TMA shape/stride for O. + template + static auto makeTmaShapeStrideO(FmhaOptions const& options) { + // + // TODO: refactor this as makeTmaShapeStrideQ when removing cutlass tma copy. + // + + // The number of tokens. + int32_t numTokens{options.mSumOfSeqLensQ}; + + // The number of heads per K/V head. + int32_t numHeadsQPerKv{options.mNumHeadsQPerKv}; + + // The batch dimension. + int32_t batchSize{1}; + + // The cute tensor shape for Q/O: (numTokens, headDim, ((numHeadsKv, numHeadsQPerKv), + // batchSize)). This maps to flattened TMA shape for Q/O: (headDim, numTokens, numHeadsKv. + // numHeadsQPerKv, batchSize). Note that TMA descriptor expects the first dimension's stride to + // be 1, so swap the first two dimension so that the headDim dimension comes first. + auto shape = std::vector{ + static_cast(options.mHeadDimV), static_cast(numTokens), + static_cast(options.mNumHeadsKv), static_cast(numHeadsQPerKv), + static_cast(batchSize)}; + + // The hidden dimension. + int32_t const hiddenDimO{options.mNumHeadsQ * options.mHeadDimV}; + + // The stride between tokens. + int32_t strideTokens{hiddenDimO}; + + // The stride between Q heads. + int32_t strideHeadsQ{options.mNumHeadsKv * options.mHeadDimV}; + + // The stride between sequences. + int32_t strideBatch{0}; + + // The stride in between K/V heads. + int32_t strideHeadsKv{options.mHeadDimV}; + // Assemble the stride (strideTokens, 1, ((strideHeadsKv, strideHeadsQ), strideBatch)). + // Swap the first two dimension as mentioned before. + auto stride = std::vector{ + 1, static_cast(strideTokens), static_cast(strideHeadsKv), + static_cast(strideHeadsQ), static_cast(strideBatch)}; + + return std::make_tuple(shape, stride); + } + + // Create the shape for K and V. + template + static auto makeShapeKv(FmhaOptions const& options, KernelParams const& params) { + // The number of keys/vals. WARNING: The if/else-if are sorted by priority. + int32_t numKeysVals{options.mMaxSeqLenKv}; + if (isPagedKv(options.mQkvLayout)) { + numKeysVals = options.mNumTokensPerPage; + } else if (isContiguousKv(options.mQkvLayout)) { + numKeysVals = options.mMaxSeqLenCacheKv; + } else { + numKeysVals = options.mSumOfSeqLensKv; + } + + // The number of heads per K/V head (packed in the sequence length for mGroupsHeadsQ). + int32_t numHeadsKv{options.mNumHeadsKv}; + + // The batch dimension. WARNING: The if/else-if are sorted by priority. + int32_t batchSize{options.mBatchSize}; + if (isPagedKv(options.mQkvLayout)) { + batchSize = params.mNumPagesInMemPool; + } else if (isContiguousKv(options.mQkvLayout)) { + batchSize = options.mBatchSize; + } else { + batchSize = 1; + } + + // Return the number of keys and batch. + return std::make_tuple(numKeysVals, numHeadsKv, batchSize); + } + + // Compute the strides for K and V. + template + static auto makeStrideKv(FmhaOptions const& options, bool isK) { + int strideKeysVals = options.kvStrideKeysValues; + int strideHeads = options.kvStrideHeads; + int strideBatch = options.kvStrideBatch; + + // The 3 strides (the other ones are 1 and 0). + return std::make_tuple(strideKeysVals, strideHeads, strideBatch); + } + + // Create the TMA shape/stride for K. + template + static auto makeTmaShapeStrideKv(FmhaOptions const& options, KernelParams const& params, + Data_type dtypeKv, bool isK) { + // The shape elements. + auto [numKeys, numHeadsQPerKv, batchSize] = makeShapeKv(options, params); + // The stride elements. + auto [strideKeys, strideHeads, strideBatch] = makeStrideKv(options, isK); + + // The headDim. + // Note that contiguousKv or pagedKv will pad K and V to maxHeadDimKv. + int32_t headDim = isK ? options.mHeadDimQk : options.mHeadDimV; + if (isPagedKv(options.mQkvLayout) || isContiguousKv(options.mQkvLayout)) { + headDim = std::max(options.mHeadDimQk, options.mHeadDimV); + } + + // For K, the cute layout: (numKeys, headDim, ((numHeadsQPerKv, numHeadsKv), + // batchSize)):(strideKeys, _1, _0, strideHeads, strideBatch). Cute swaps the first two + // dimension (to make sure stride of first dimension is 1) and ignores the numHeadsQPerKv + // dimension (it's stride is always 0). For V, the headDim dimension is already the first + // dimension so no swapping is needed. + + // Therefore, the resulting TMA layout is 4D: (headDim, numKeys, numHeadsKv, batchSize):(1, + // strideKeys, strideHeads, strideBatch) + + // Note that for FP4 KV input, elements are stored as uint8_t, each packs 2 FP4 elements. + // The column index and strides needs to divide by 2. + auto const colIdxDivisor = dtypeKv == DATA_TYPE_E2M1 ? 2 : 1; + auto shape = std::vector{ + static_cast(headDim / colIdxDivisor), static_cast(numKeys), + static_cast(options.mNumHeadsKv), static_cast(batchSize)}; + auto stride = std::vector{1, static_cast(strideKeys / colIdxDivisor), + static_cast(strideHeads / colIdxDivisor), + static_cast(strideBatch / colIdxDivisor)}; + + return std::make_tuple(shape, stride); + } + + // Create the TMA shape/stride for KV scaling factors. + template + static auto makeTmaShapeStrideKvSf(FmhaOptions const& options, KernelParams const& params, + bool isK) { + // The shape elements. + auto [numKeys, numHeadsQPerKv, batchSize] = makeShapeKv(options, params); + // The stride elements. + auto [strideKeys, strideHeads, strideBatch] = makeStrideKv(options, isK); + + // The headDim. + // Note that contiguousKv or pagedKv will pad K and V to maxHeadDimKv. + int32_t headDim = isK ? options.mHeadDimQk : options.mHeadDimV; + if (isPagedKv(options.mQkvLayout) || isContiguousKv(options.mQkvLayout)) { + headDim = std::max(options.mHeadDimQk, options.mHeadDimV); + } + // The number of elements per SF. + int32_t NumEltsPerSf = 16; + + // The KV shape is: (headDim, numKeys, numHeadsKv, batchSize) + // Therefore, the KV SF shape should be (headDim / NumEltsPerSf, numKeys, numHeadsKv, + // batchSize). Considering the TMA requires box width to be multiple of 16B, without changing + // the underlying layout, we reshape into (16, numKeys * headDim / NumEltsPerSf / 16, + // numHeadsKv, batchSize) + + // Note that it only works for pagedKv layout. + TORCH_CHECK(isPagedKv(options.mQkvLayout), "The qkvLayout is not supported."); + + auto shape = std::vector{ + 16, static_cast(numKeys * headDim / NumEltsPerSf / 16), + static_cast(options.mNumHeadsKv), static_cast(batchSize)}; + auto stride = std::vector{1, 16, static_cast(strideHeads / NumEltsPerSf), + static_cast(strideBatch / NumEltsPerSf)}; + + return std::make_tuple(shape, stride); + } + + // Prepare pointers for TMA descriptors. + static std::tuple getDevicePtrs( + TllmGenFmhaRunnerParams const& runnerParams, int32_t bytesPerElt) { + // Declare the q, k, v ptrs. + void const *qPtr{runnerParams.qPtr}, *kPtr{runnerParams.kPtr}, *vPtr{runnerParams.vPtr}; + + // Set Q, K and V pointer from packed QKV tensor. + if (isPackedQkv(runnerParams.mQkvLayout)) { + qPtr = runnerParams.qkvPtr; + kPtr = reinterpret_cast(reinterpret_cast(runnerParams.qkvPtr) + + runnerParams.mNumHeadsQ * runnerParams.mHeadDimQk * + bytesPerElt); + vPtr = reinterpret_cast(reinterpret_cast(runnerParams.qkvPtr) + + (runnerParams.mNumHeadsQ + runnerParams.mNumHeadsKv) * + runnerParams.mHeadDimQk * bytesPerElt); + } + // Set K and V pointer from pagedKv tensor. + else if (isPagedKv(runnerParams.mQkvLayout)) { + // Note that the offsets will be fully handled by the pageIdx buffer. + kPtr = runnerParams.kPtr; + vPtr = runnerParams.vPtr; + } + // Set K and V pointer from contiguousQAnddKv tensor. + else if (isContiguousKv(runnerParams.mQkvLayout)) { + kPtr = runnerParams.kvPtr; + // The maximum headDim of K and V. + // Note that contiguousKv or pagedKv will pad K and V to maxHeadDimKv. + int32_t const maxHeadDimKv{std::max(runnerParams.mHeadDimQk, runnerParams.mHeadDimV)}; + vPtr = reinterpret_cast( + reinterpret_cast(runnerParams.kvPtr) + + runnerParams.mNumHeadsKv * runnerParams.mMaxSeqLenCacheKv * maxHeadDimKv * bytesPerElt); + } + + // Return the pointers. + return std::make_tuple(qPtr, kPtr, vPtr); + } + + // Build tma descriptors. + template + static CUtensorMap buildNdTmaDescriptor(FmhaOptions const& options, Data_type dtypeElt, + std::vector const& shapes, + std::vector const& strides, + std::vector const& tileShapes, void* gmemAddr, + bool swizzled = true) { + CUtensorMap desc{}; + // The data type. + CUtensorMapDataType tmaDataFormat; + if (dtypeElt == DATA_TYPE_E2M1 || dtypeElt == DATA_TYPE_E4M3) { + tmaDataFormat = CU_TENSOR_MAP_DATA_TYPE_UINT8; + } else if (dtypeElt == DATA_TYPE_FP16) { + tmaDataFormat = CU_TENSOR_MAP_DATA_TYPE_FLOAT16; + } else if (dtypeElt == DATA_TYPE_BF16) { + tmaDataFormat = CU_TENSOR_MAP_DATA_TYPE_BFLOAT16; + } else { + TORCH_CHECK(false, "Unexpected dtype %d", static_cast(dtypeElt)); + } + + // The swizzle type. + CUtensorMapSwizzle swizzleType; + int32_t numBytesInLeadingDim = tileShapes[0] * get_size_in_bits(dtypeElt) / 8 /*bits*/; + if (!swizzled) { + swizzleType = CU_TENSOR_MAP_SWIZZLE_NONE; + } else if ((numBytesInLeadingDim % 128) == 0) { + swizzleType = CU_TENSOR_MAP_SWIZZLE_128B; + } else if ((numBytesInLeadingDim % 64) == 0) { + swizzleType = CU_TENSOR_MAP_SWIZZLE_64B; + } else if ((numBytesInLeadingDim % 32) == 0) { + swizzleType = CU_TENSOR_MAP_SWIZZLE_32B; + } else { + TORCH_CHECK(false, "Unexpected numBytesInLeadingDim %d", numBytesInLeadingDim); + } + + // Check gmem address must be 16B-aligned + TORCH_CHECK((reinterpret_cast(gmemAddr) & 0b1111) == 0); + + // Check shape must be in range [1, 2^32] + int32_t dim = shapes.size(); + // Max five dimension and min 3 dimension. + TORCH_CHECK((dim <= 5) && (dim >= 3)); + // Check shape range. + for (int32_t ii = 0; ii < dim; ++ii) { + TORCH_CHECK(shapes[ii] >= (uint64_t(1))); // Size must be min 1 + TORCH_CHECK(shapes[ii] <= (uint64_t(1) << 32)); // Size must be max 2^32 + } + + // TMA descriptor does not store the zeroth stride and assumes it is 1. + TORCH_CHECK(static_cast(strides.size()) == dim); + TORCH_CHECK(strides[0] == 1); + + // Build strides in bytes. + // cuTensorMapEncodeTiled ignores the stride of the first dimension (implicitly 1). + std::vector stridesInBytes(dim - 1); + for (int32_t ii = 0; ii < dim - 1; ++ii) { + stridesInBytes[ii] = strides[ii + 1] * + std::max(get_size_in_bits(dtypeElt), static_cast(8)) / 8 /*bit*/; + } + + // Set tile strides to 0; + std::vector tileStrides(dim, 1); + + // Build the descriptor. + CUresult result = + cuTensorMapEncodeTiled(&desc, tmaDataFormat, + /*tensorRank=*/dim, gmemAddr, shapes.data(), stridesInBytes.data(), + tileShapes.data(), tileStrides.data(), + /*interleave=*/CU_TENSOR_MAP_INTERLEAVE_NONE, swizzleType, + /*l2Promotion=*/CU_TENSOR_MAP_L2_PROMOTION_L2_128B, + /*oobFill=*/CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + + if (result != CUDA_SUCCESS) { + char const* err_str; + cuGetErrorString(result, &err_str); + std::cerr << "Error: Failed to initialize the TMA descriptor due to " << err_str << std::endl; + std::cerr << "tmaFormat: " << static_cast(tmaDataFormat) << " dim: " << dim + << " gmem: " << gmemAddr << std::endl; + std::cerr << "Shape: " << shapes[0] << " " << shapes[1] << " " << shapes[2] << " " + << shapes[3] << " " << shapes[4] << std::endl; + std::cerr << "Stride: " << stridesInBytes[0] << " " << stridesInBytes[1] << " " + << stridesInBytes[2] << " " << stridesInBytes[3] << std::endl; + std::cerr << "tileShapes: " << tileShapes[0] << " " << tileShapes[1] << " " << tileShapes[2] + << " " << tileShapes[3] << " " << tileShapes[4] << std::endl; + std::cerr << "tileStrides: " << tileStrides[0] << " " << tileStrides[1] << " " + << tileStrides[2] << " " << tileStrides[3] << " " << tileStrides[4] << std::endl; + std::cerr << "swizzleType: " << int(swizzleType) << std::endl; + TORCH_CHECK(false); + } + + return desc; + } + + // Setup the kernel parameters. + template + static KernelParams setKernelParams(FmhaOptions_ const& options, KernelMeta const& kernelMeta, + int32_t maxNumCtasQ, int32_t maxNumCtasKv) { + // Create the return struct. + KernelParams params; + + // Get the device pointers for TMA descriptors. + auto [qPtr, kPtr, vPtr] = getDevicePtrs(options, get_size_in_bytes(kernelMeta.mDataTypeKv)); + + // The maximum headDim of K and V. + // Note that contiguousKv or pagedKv will pad K and V to maxHeadDimKv. + int32_t const maxHeadDimKv{std::max(options.mHeadDimQk, options.mHeadDimV)}; + + // Set the number of pages in the memory pool for paged K/V cache. + if (isPagedKv(options.mQkvLayout)) { + params.mNumPagesInMemPool = options.mNumPagesInMemPool == 0 + ? options.mMaxNumPagesPerSeqKv * 2 * options.mBatchSize + : options.mNumPagesInMemPool; + } + + // The number of elements in 128B for Q. + int32_t numEltsIn128BQ = (128 * 8) / get_size_in_bits(kernelMeta.mDataTypeQ); + // The number of head elts (per token) in each block of shared memory. + int32_t numEltsInClampedHeadDimQ = std::min(numEltsIn128BQ, options.mHeadDimQk); + + // Shape/stride for gmem tensor Q. + auto [shapeQ, strideQ, tileShapeQ] = makeTmaShapeStrideQ( + options, kernelMeta.mGroupsHeadsQ, kernelMeta.mTileSizeQ, numEltsInClampedHeadDimQ); + // Build tma descriptor for Q. + params.tmaQ_ = buildNdTmaDescriptor(options, kernelMeta.mDataTypeQ, shapeQ, strideQ, tileShapeQ, + const_cast(qPtr)); + + // The number of keys per tile. + int32_t numKeysPerTile = isPagedKv(options.mQkvLayout) + ? std::min(options.mNumTokensPerPage, kernelMeta.mTileSizeKv) + : kernelMeta.mTileSizeKv; + // The number of elements in 128B for Q. + int32_t numEltsIn128BKv = (128 * 8) / get_size_in_bits(kernelMeta.mDataTypeKv); + // The number of head elts (per token) in each block of shared memory (see above explanation). + int32_t numEltsInClampedHeadDimKv = std::min(numEltsIn128BKv, maxHeadDimKv); + + // Shape/stride for gmem tensor Kv. + auto [shapeK, strideK] = + makeTmaShapeStrideKv(options, params, kernelMeta.mDataTypeKv, /*isK*/ true); + auto [shapeV, strideV] = + makeTmaShapeStrideKv(options, params, kernelMeta.mDataTypeKv, /*isK*/ false); + // Build tma descriptor for K. + // Do we have to transform K/V before MMA? + bool const transformsKv{kernelMeta.mDataTypeKv != kernelMeta.mDataTypeQ}; + // Note that for FP4 KV input, elements are stored as uint8_t, each packs 2 FP4 elements. + auto const numEltsDivisor = kernelMeta.mDataTypeKv == DATA_TYPE_E2M1 ? 2 : 1; + // The tileShapes for K/V. + std::vector tileShapeKv(shapeK.size(), 1); + tileShapeKv[0] = numEltsInClampedHeadDimKv / numEltsDivisor; + tileShapeKv[1] = numKeysPerTile; + // Build tma descriptor for K. + params.tmaK_ = buildNdTmaDescriptor(options, kernelMeta.mDataTypeKv, shapeK, strideK, + tileShapeKv, const_cast(kPtr), + /*swizzled = */ !transformsKv); + // Build tma descriptor for V. + params.tmaV_ = buildNdTmaDescriptor(options, kernelMeta.mDataTypeKv, shapeV, strideV, + tileShapeKv, const_cast(vPtr), + /*swizzled = */ !transformsKv); + + // If the KV dtype is E2m1, additional scaling factors are needed for dequant. + if (kernelMeta.mDataTypeKv == DATA_TYPE_E2M1) { + // The number of elements per SF. + int32_t NumEltsPerSf = 16; + // Compute the shape and stride for SF tensor. + // FIXME: assume K and V uses the same shape. + auto [shapeKvSf, strideKvSf] = makeTmaShapeStrideKvSf(options, params, /*isK*/ true); + + // The tileShapes for K/V. + std::vector tileShapeKvSf(shapeKvSf.size(), 1); + tileShapeKvSf[0] = 16; + tileShapeKvSf[1] = numKeysPerTile * maxHeadDimKv / NumEltsPerSf / 16; + + // The tile box is reshaped from (headDim / NumEltsPerSf, tileSizeKv) into (16, tileSizeKv * + // headDim / NumEltsPerSf / 16). See makeTmaShapeStrideKvSf for details. Build tma descriptor + // for K SF. + params.tmaKSf_ = buildNdTmaDescriptor(options, DATA_TYPE_E4M3, shapeKvSf, strideKvSf, + tileShapeKvSf, const_cast(options.kSfBasePtr), + /*swizzled = */ false); + + // Build tma descriptor for V SF. + params.tmaVSf_ = buildNdTmaDescriptor(options, DATA_TYPE_E4M3, shapeKvSf, strideKvSf, + tileShapeKvSf, const_cast(options.vSfBasePtr), + /*swizzled = */ false); + } + + // Shape/stride for gmem tensor O. + auto [shapeO, strideO] = makeTmaShapeStrideO(options); + // The tileShapes for O. + std::vector tileShapeO(shapeO.size(), 1); + tileShapeO[0] = numEltsInClampedHeadDimQ; + tileShapeO[1] = kernelMeta.mTileSizeQ; + // Build tma descriptor for O. + params.tmaO_ = buildNdTmaDescriptor(options, kernelMeta.mDataTypeQ, shapeO, strideO, tileShapeO, + const_cast(options.oPtr)); + + // Set the other kernel parameters. + params.ptrCumSeqLensQ = options.cumSeqLensQPtr; + params.ptrCumSeqLensKv = options.cumSeqLensKvPtr; + + // The packed custom mask. + params.ptrCustomMask = options.customMaskPtr; + // The packed custom mask's offsets of each sequence. + params.ptrCustomMaskOffsets = options.customMaskOffsetsPtr; + // The first sparseMask offsets in the Kv sequence dimension. + params.ptrFirstSparseMaskOffsetsKv = options.firstSparseMaskOffsetsKvPtr; + + // The output buffer. + params.ptrO = options.oPtr; + // The output scaling factor buffer. + params.ptrSfO = options.oSfPtr; + + // TRT-LLM restrictions: the quantization scales must be on the device. + params.ptrOutputScale = options.outputScalePtr; + + // The sequence lengths for Kv. + params.ptrSeqLensKv = options.seqLensKvPtr; + + // Attention sink + params.ptrAttentionSinks = options.ptrAttentionSinks; + + // The partial buffers' pointers when the multiCtasKv mode is enabled. + int64_t partialStatsBufferSize = options.mMultiProcessorCount * kernelMeta.mStepQ; + params.ptrMultiCtasKvCounter = options.multiCtasKvCounterPtr; + params.ptrPartialStats = reinterpret_cast(options.multiCtasKvScratchPtr); + params.ptrPartialO = params.ptrPartialStats + partialStatsBufferSize; + + params.ptrPageIdxKv = options.kvPageIdxPtr; + params.ptrScaleSoftmaxLog2 = options.scaleSoftmaxLog2Ptr; + + params.ptrScaleSfKv = options.kvSfScalePtr; + params.ptrScaleSfO = options.oSfScalePtr; + params.mScaleSfO = options.mScaleSfO; + + params.mAttentionWindowSize = options.mAttentionWindowSize; + if (isSlidingOrChunkedCausalMask( + static_cast(kernelMeta.mMaskType)) && + options.mChunkedAttentionSize != INT_MAX) { + TORCH_CHECK((options.mChunkedAttentionSize & (options.mChunkedAttentionSize - 1)) == 0, + "Chunked attention size must be a power of 2"); + params.mChunkedAttentionSizeLog2 = std::log2(options.mChunkedAttentionSize); + } else { + // Default 0 means that chunked attention is disabled. + params.mChunkedAttentionSizeLog2 = 0; + } + params.mMaxSeqLenQ = options.mMaxSeqLenQ; + params.mMaxSeqLenKv = options.mMaxSeqLenKv; + params.mMaxNumCtasQ = maxNumCtasQ; + params.mMaxNumCtasKv = maxNumCtasKv; + params.mMaxNumPagesPerSeqKv = options.mMaxNumPagesPerSeqKv; + // TODO: just use mMaxSeqLenQ for number of MTP tokens. + params.mNumMtpTokens = options.mMaxSeqLenQ; + params.mSumOfSeqLensQ = options.mSumOfSeqLensQ; + params.mSumOfSeqLensKv = options.mSumOfSeqLensKv; + params.mBatchSize = options.mBatchSize; + params.mChunkedAttentionSizeLog2 = 0; + params.mNumHeadsQ = options.mNumHeadsQ; + params.mNumHeadsKv = options.mNumHeadsKv; + params.mNumHeadsQPerKv = options.mNumHeadsQPerKv; + params.mNumHiddenEltsO = options.mNumHeadsQ * options.mHeadDimQk; + // todo(Yingyi): might take a scalar tensor later + params.mOutputScale = options.outputScale; + params.mScaleSoftmaxLog2 = options.scaleSoftmaxLog2; + params.mStartTokenIdxSfO = options.mSfStartTokenIdx; + params.mScaleSfKv = options.mScaleSfKv; + params.ptrSoftmaxStats = nullptr; + return params; + } +}; diff --git a/include/flashinfer/trtllm/fused_moe/DevKernel.h b/include/flashinfer/trtllm/fused_moe/DevKernel.h new file mode 100644 index 0000000000000000000000000000000000000000..15e636d39bc1c05fb60bab4a7ee81933f5bd57a5 --- /dev/null +++ b/include/flashinfer/trtllm/fused_moe/DevKernel.h @@ -0,0 +1,421 @@ +/* + * Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +#include +#include + +#include "flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/DtypeDecl.h" +#include "flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/SfLayoutDecl.h" +// #include +#include +#include +#include +#include +// #include +#include "flashinfer/trtllm/common/cudaUtils.h" + +namespace moe::dev { + +#define CHECK_CUDA_ERROR(cmd) \ + do { \ + cudaError_t e = cmd; \ + if (e != cudaSuccess) { \ + std::cout << "CUDA error in " << __FILE__ << ":" << __LINE__ << " executing '" << #cmd \ + << "': " << cudaGetErrorString(e); \ + } \ + TORCH_CHECK(e == cudaSuccess, "Got CUDA error. See above for details."); \ + } while (0) +//////////////////////////////////////////////////////////////////////////////////////////////////// + +#define LAUNCH_ESC(...) __VA_ARGS__ + +#define LAUNCH_PDL(data, coopLaunch, types, kernel, numBlocks, numThreads, smemSize, stream) \ + cudaLaunchConfig_t config{}; \ + config.gridDim = numBlocks; \ + config.blockDim = numThreads; \ + config.dynamicSmemBytes = smemSize; \ + config.stream = (cudaStream_t)stream; \ + \ + cudaLaunchAttribute attributes[2] = {}; \ + attributes[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; \ + attributes[0].val.programmaticStreamSerializationAllowed = int(data.mUsePdl); \ + attributes[1].id = cudaLaunchAttributeCooperative; \ + attributes[1].val.cooperative = int(coopLaunch); \ + config.attrs = attributes; \ + config.numAttrs = 2; \ + if (data.mUsePdl) { \ + auto params = KernelParams::setKernelParams(data); \ + auto kernelTyped = kernel>; \ + if (smemSize > 48 * 1024) \ + CHECK_CUDA_ERROR(cudaFuncSetAttribute( \ + kernelTyped, cudaFuncAttributeMaxDynamicSharedMemorySize, smemSize)); \ + CHECK_CUDA_ERROR(cudaLaunchKernelEx(&config, kernelTyped, params)); \ + } else { \ + auto params = KernelParams::setKernelParams(data); \ + auto kernelTyped = kernel>; \ + if (smemSize > 48 * 1024) \ + CHECK_CUDA_ERROR(cudaFuncSetAttribute( \ + kernelTyped, cudaFuncAttributeMaxDynamicSharedMemorySize, smemSize)); \ + CHECK_CUDA_ERROR(cudaLaunchKernelEx(&config, kernelTyped, params)); \ + } + +#define LAUNCH(data, kernel, numBlocks, numThreads, smemSize, stream) \ + if (data.mDtypeElt == tg::Dtype::Fp16) { \ + LAUNCH_PDL(data, false, cutlass::half_t, kernel, numBlocks, numThreads, smemSize, stream); \ + } else if (data.mDtypeElt == tg::Dtype::E4m3) { \ + LAUNCH_PDL(data, false, cutlass::float_e4m3_t, kernel, numBlocks, numThreads, smemSize, \ + stream); \ + } else if (data.mDtypeElt == tg::Dtype::Bfloat16) { \ + LAUNCH_PDL(data, false, cutlass::bfloat16_t, kernel, numBlocks, numThreads, smemSize, stream); \ + } else { \ + TORCH_WARN("Unsupported dtypeElt"); \ + } + +#define LAUNCH_EXPW(data, kernel, numBlocks, numThreads, smemSize, stream) \ + if (data.mDtypeElt == tg::Dtype::Fp16 && data.mDtypeExpW == tg::Dtype::Fp32) { \ + LAUNCH_PDL(data, false, LAUNCH_ESC(cutlass::half_t, float), kernel, numBlocks, numThreads, \ + smemSize, stream); \ + } else if (data.mDtypeElt == tg::Dtype::E4m3 && data.mDtypeExpW == tg::Dtype::Fp32) { \ + LAUNCH_PDL(data, false, LAUNCH_ESC(cutlass::float_e4m3_t, float), kernel, numBlocks, \ + numThreads, smemSize, stream); \ + } else if (data.mDtypeElt == tg::Dtype::Bfloat16 && data.mDtypeExpW == tg::Dtype::Fp32) { \ + LAUNCH_PDL(data, false, LAUNCH_ESC(cutlass::bfloat16_t, float), kernel, numBlocks, numThreads, \ + smemSize, stream); \ + } else if (data.mDtypeElt == tg::Dtype::Fp16 && data.mDtypeExpW == tg::Dtype::Bfloat16) { \ + LAUNCH_PDL(data, false, LAUNCH_ESC(cutlass::half_t, cutlass::bfloat16_t), kernel, numBlocks, \ + numThreads, smemSize, stream); \ + } else if (data.mDtypeElt == tg::Dtype::E4m3 && data.mDtypeExpW == tg::Dtype::Bfloat16) { \ + LAUNCH_PDL(data, false, LAUNCH_ESC(cutlass::float_e4m3_t, cutlass::bfloat16_t), kernel, \ + numBlocks, numThreads, smemSize, stream); \ + } else if (data.mDtypeElt == tg::Dtype::Bfloat16 && data.mDtypeExpW == tg::Dtype::Bfloat16) { \ + LAUNCH_PDL(data, false, LAUNCH_ESC(cutlass::bfloat16_t, cutlass::bfloat16_t), kernel, \ + numBlocks, numThreads, smemSize, stream); \ + } else { \ + TORCH_WARN("Unsupported pair"); \ + } + +#define LAUNCH_ROUTING(data, coopLaunch, kernel, numBlocks, numThreads, smemSize, stream) \ + if (data.mDtypeExpW == tg::Dtype::Fp32) { \ + LAUNCH_PDL(data, coopLaunch, LAUNCH_ESC(float, float), kernel, numBlocks, numThreads, \ + smemSize, stream); \ + } else if (data.mDtypeExpW == tg::Dtype::Bfloat16) { \ + LAUNCH_PDL(data, coopLaunch, LAUNCH_ESC(__nv_bfloat16, __nv_bfloat16), kernel, numBlocks, \ + numThreads, smemSize, stream); \ + } else { \ + TORCH_WARN("Unsupported dtypeExpW"); \ + } + +#define LAUNCH_ROUTING_WITH_EXTRA_FLAG(data, coopLaunch, kernel, numBlocks, numThreads, smemSize, \ + stream, extraFlag, forceFloatInput) \ + if (data.mDtypeExpW == tg::Dtype::Fp32 && extraFlag) { \ + LAUNCH_PDL(data, coopLaunch, LAUNCH_ESC(float, float, true), kernel, numBlocks, numThreads, \ + smemSize, stream); \ + } else if (data.mDtypeExpW == tg::Dtype::Fp32) { \ + LAUNCH_PDL(data, coopLaunch, LAUNCH_ESC(float, float, false), kernel, numBlocks, numThreads, \ + smemSize, stream); \ + } else if (data.mDtypeExpW == tg::Dtype::Bfloat16 && extraFlag && forceFloatInput) { \ + LAUNCH_PDL(data, coopLaunch, LAUNCH_ESC(float, __nv_bfloat16, true), kernel, numBlocks, \ + numThreads, smemSize, stream); \ + } else if (data.mDtypeExpW == tg::Dtype::Bfloat16 && extraFlag) { \ + LAUNCH_PDL(data, coopLaunch, LAUNCH_ESC(__nv_bfloat16, __nv_bfloat16, true), kernel, \ + numBlocks, numThreads, smemSize, stream); \ + } else if (data.mDtypeExpW == tg::Dtype::Bfloat16 && forceFloatInput) { \ + LAUNCH_PDL(data, coopLaunch, LAUNCH_ESC(float, __nv_bfloat16, false), kernel, numBlocks, \ + numThreads, smemSize, stream); \ + } else if (data.mDtypeExpW == tg::Dtype::Bfloat16) { \ + LAUNCH_PDL(data, coopLaunch, LAUNCH_ESC(__nv_bfloat16, __nv_bfloat16, false), kernel, \ + numBlocks, numThreads, smemSize, stream); \ + } else { \ + TORCH_WARN("Unsupported dtypeExpW"); \ + } + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace activation { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace tg = batchedGemm::trtllm::gen; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +struct Data { + tg::Dtype mDtypeElt{tg::Dtype::Fp16}; + bool mUsePdl{false}; + bool mUseDeepSeekFp8{false}; + + void* inPtr; + void* outPtr; + float* inDqSfsPtr = nullptr; + float* outDqSfsPtr = nullptr; + + int32_t innerDim; + int32_t numTokens; + int32_t topK; + int32_t* expandedIdxToPermutedIdx; + + int32_t const* totalNumPaddedTokens; +}; + +template +struct KernelParams { + using Type = Type_; + static constexpr bool UsePdl = UsePdl_; + + Type const* inPtr; + Type* outPtr; + + float* inDqSfsPtr = nullptr; + float* outDqSfsPtr = nullptr; + + int32_t innerDim; + int32_t numTokens; + int32_t topK; + int32_t* expandedIdxToPermutedIdx; + + int32_t const* totalNumPaddedTokens; + + static KernelParams setKernelParams(Data const& data) { + KernelParams params; + + params.inPtr = (Type*)data.inPtr; + params.outPtr = (Type*)data.outPtr; + params.inDqSfsPtr = data.inDqSfsPtr; + params.outDqSfsPtr = data.outDqSfsPtr; + + params.expandedIdxToPermutedIdx = data.expandedIdxToPermutedIdx; + + params.innerDim = data.innerDim; + params.numTokens = data.numTokens; + params.topK = data.topK; + params.totalNumPaddedTokens = data.totalNumPaddedTokens; + + return params; + } +}; + +void run(Data const& data, void* stream); + +} // namespace activation + +namespace convertsf { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace tg = batchedGemm::trtllm::gen; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +struct Data { + bool mUsePdl{false}; + + void* inSfPtr = nullptr; + void* outSfPtr = nullptr; + int32_t hiddenDimSf; + int32_t numTokens; + tg::SfLayout sfLayoutSrc; + tg::SfLayout sfLayoutDst; +}; + +template +struct KernelParams { + using Type = Type_; + static constexpr bool UsePdl = UsePdl_; + + void const* inSfPtr = nullptr; + void* outSfPtr = nullptr; + int32_t hiddenDimSf; + int32_t numTokens; + tg::SfLayout sfLayoutSrc; + tg::SfLayout sfLayoutDst; + + static KernelParams setKernelParams(Data const& data) { + KernelParams params; + + params.inSfPtr = data.inSfPtr; + params.outSfPtr = data.outSfPtr; + params.hiddenDimSf = data.hiddenDimSf; + params.numTokens = data.numTokens; + params.sfLayoutSrc = data.sfLayoutSrc; + params.sfLayoutDst = data.sfLayoutDst; + + return params; + } +}; + +void run(Data const& data, void* stream); + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace convertsf + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace permute { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace tg = batchedGemm::trtllm::gen; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +struct Data { + tg::Dtype mDtypeElt{tg::Dtype::Fp16}; + bool mUsePdl{false}; + bool mUseDeepSeekFp8{false}; + + void* inPtr; + void* outPtr; + float* inDqSfsPtr = nullptr; + float* outDqSfsPtr = nullptr; + int32_t* expandedIdxToPermutedIdx; + int32_t hiddenDim; + int32_t numTokens; + int32_t topK; + int32_t const* totalNumPaddedTokens; +}; + +template +struct KernelParams { + using Type = Type_; + static constexpr bool UsePdl = UsePdl_; + + Type const* inPtr; + Type* outPtr; + float const* inDqSfsPtr; + float* outDqSfsPtr; + int32_t* expandedIdxToPermutedIdx; + int32_t hiddenDim; + int32_t numTokens; + int32_t topK; + int32_t const* totalNumPaddedTokens; + bool useDeepSeekFp8; + + static KernelParams setKernelParams(Data const& data) { + KernelParams params; + + params.inPtr = (Type*)data.inPtr; + params.outPtr = (Type*)data.outPtr; + params.inDqSfsPtr = data.inDqSfsPtr; + params.outDqSfsPtr = data.outDqSfsPtr; + params.expandedIdxToPermutedIdx = data.expandedIdxToPermutedIdx; + params.hiddenDim = data.hiddenDim; + params.numTokens = data.numTokens; + params.topK = data.topK; + params.totalNumPaddedTokens = data.totalNumPaddedTokens; + params.useDeepSeekFp8 = data.mUseDeepSeekFp8; + + return params; + } +}; + +void run(Data const& data, void* stream); + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace permute + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace finalize { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace tg = batchedGemm::trtllm::gen; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +struct Data { + tg::Dtype mDtypeElt{tg::Dtype::Fp16}; + tg::Dtype mDtypeExpW{tg::Dtype::Bfloat16}; + bool mUsePdl{false}; + bool mUseDeepSeekFp8{false}; + + void* inPtr; + void* outPtr; + float* inDqSfsPtr = nullptr; + float* outDqSfsPtr = nullptr; + + void* expertWeightsPtr; + int32_t* expandedIdxToPermutedIdx; + + int32_t numTokens; + int32_t numExperts; + int32_t topK; + // Hidden dimension output of MoE block. It is not padded. + int32_t hiddenDim; + // Hidden dimension output of FC2. It might be padded. + int32_t hiddenDimPadded; + int32_t const* totalNumPaddedTokens; +}; + +template +struct KernelParams { + using Type = Type_; + using TypeExpW = TypeExpW_; + static constexpr bool UsePdl = UsePdl_; + + Type const* inPtr; + TypeExpW const* expertWeightsPtr; + Type* outPtr; + + float* inDqSfsPtr = nullptr; + float* outDqSfsPtr = nullptr; + + int32_t* expandedIdxToPermutedIdx; + + int32_t hiddenDim; + int32_t hiddenDimPadded; + int32_t numTokens; + int32_t numExperts; + int32_t topK; + int32_t const* totalNumPaddedTokens; + + static KernelParams setKernelParams(Data const& data) { + KernelParams params; + + params.inPtr = (Type*)data.inPtr; + params.expertWeightsPtr = (TypeExpW*)data.expertWeightsPtr; + params.outPtr = (Type*)data.outPtr; + params.inDqSfsPtr = data.inDqSfsPtr; + params.outDqSfsPtr = data.outDqSfsPtr; + + params.expandedIdxToPermutedIdx = data.expandedIdxToPermutedIdx; + + params.hiddenDim = data.hiddenDim; + params.hiddenDimPadded = data.hiddenDimPadded; + params.numTokens = data.numTokens; + params.numExperts = data.numExperts; + params.topK = data.topK; + params.totalNumPaddedTokens = data.totalNumPaddedTokens; + + return params; + } +}; + +void run(Data const& data, void* stream); + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace finalize + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace moe::dev diff --git a/include/flashinfer/trtllm/fused_moe/IntFastDiv.h b/include/flashinfer/trtllm/fused_moe/IntFastDiv.h new file mode 100644 index 0000000000000000000000000000000000000000..5df34943b2d6b0eac9bd77de7161aecd94a301ed --- /dev/null +++ b/include/flashinfer/trtllm/fused_moe/IntFastDiv.h @@ -0,0 +1,149 @@ +/* + * Copyright (c) 2022-2025, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +// code in this file based on https://github.com/milakov/int_fastdiv +// note: above repo is not added as a submodule because we slightly update APIs + +#include +#include + +namespace trtllm::dev { + +//////////////////////////////////////////////////////////////////////////////////////////////////// +// ************************************************************************************************* +// IntFastDiv class. +// ************************************************************************************************* +//////////////////////////////////////////////////////////////////////////////////////////////////// + +class IntFastDiv { + public: + // we allow a default constructor (and over-writing the divisor with assignment) + CUTE_HOST + IntFastDiv() : mDivisor(1), mMagicM(0), mMagicS(-1), mAddSign(1) {} + + CUTE_HOST + IntFastDiv(int divisor) : mDivisor(divisor) { + if (mDivisor == 0) throw std::runtime_error("IntFastDiv: cannot divide by 0"); + updateMagicNumbers(); + } + + CUTE_HOST + IntFastDiv& operator=(int divisor) { + this->mDivisor = divisor; + if (this->mDivisor == 0) throw std::runtime_error("IntFastDiv: cannot divide by 0"); + updateMagicNumbers(); + return *this; + } + + CUTE_HOST_DEVICE + operator int() const { return mDivisor; } + + private: + int mDivisor; + int mMagicM; + int mMagicS; + int mAddSign; + + // Hacker'mMagicS Delight, Second Edition, Chapter 10, Integer Division By Constants + CUTE_HOST + void updateMagicNumbers() { + if (mDivisor == 1) { + mMagicM = 0; + mMagicS = -1; + mAddSign = 1; + return; + } else if (mDivisor == -1) { + mMagicM = 0; + mMagicS = -1; + mAddSign = -1; + return; + } + + int p; + unsigned int tmpAd, tmpAnc, delta, q1, r1, q2, r2, t; + unsigned const two31 = 0x80000000; + tmpAd = abs(mDivisor); + t = two31 + ((unsigned int)mDivisor >> 31); + tmpAnc = t - 1 - t % tmpAd; + p = 31; + q1 = two31 / tmpAnc; + r1 = two31 - q1 * tmpAnc; + q2 = two31 / tmpAd; + r2 = two31 - q2 * tmpAd; + do { + ++p; + q1 = 2 * q1; + r1 = 2 * r1; + if (r1 >= tmpAnc) { + ++q1; + r1 -= tmpAnc; + } + q2 = 2 * q2; + r2 = 2 * r2; + if (r2 >= tmpAd) { + ++q2; + r2 -= tmpAd; + } + delta = tmpAd - r2; + } while (q1 < delta || (q1 == delta && r1 == 0)); + this->mMagicM = q2 + 1; + if (mDivisor < 0) this->mMagicM = -this->mMagicM; + this->mMagicS = p - 32; + + if ((mDivisor > 0) && (mMagicM < 0)) + mAddSign = 1; + else if ((mDivisor < 0) && (mMagicM > 0)) + mAddSign = -1; + else + mAddSign = 0; + } + + CUTE_HOST_DEVICE + friend int operator/(int const dividend, IntFastDiv const& divisor); +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +CUTE_HOST_DEVICE +int operator/(int const dividend, IntFastDiv const& divisor) { + int q; +#ifdef __CUDA_ARCH__ + asm("mul.hi.s32 %0, %1, %2;" : "=r"(q) : "r"(divisor.mMagicM), "r"(dividend)); +#else + q = (((unsigned long long)((long long)divisor.mMagicM * (long long)dividend)) >> 32); +#endif + q += dividend * divisor.mAddSign; + if (divisor.mMagicS >= 0) { + q >>= divisor.mMagicS; // we rely on this to be implemented as arithmetic shift + q += (((unsigned int)q) >> 31); + } + return q; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +CUTE_HOST_DEVICE +int operator%(int const dividend, IntFastDiv const& divisor) { + int quotient = dividend / divisor; + int remainder = dividend - quotient * divisor; + return remainder; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace trtllm::dev diff --git a/include/flashinfer/trtllm/fused_moe/RoutingKernel.cuh b/include/flashinfer/trtllm/fused_moe/RoutingKernel.cuh new file mode 100644 index 0000000000000000000000000000000000000000..3422e2999adc0ec5b45bf65d16d9ef6f71ea6ca0 --- /dev/null +++ b/include/flashinfer/trtllm/fused_moe/RoutingKernel.cuh @@ -0,0 +1,684 @@ +/* + * Copyright (c) 2022-2025, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include +#include +#include + +#include +#include +#include + +#include "DevKernel.h" +#include "RoutingKernel.h" +#include "RoutingKernelTopK.cuh" + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace moe::dev { + +//////////////////////////////////////////////////////////////////////////////////////////////////// +namespace routing { + +namespace cg = cooperative_groups; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +static constexpr int WarpSize = 32; +static constexpr int NumBlocksPerCluster = 8; +// Performance tuning knob. +static constexpr int NumEltsPerOffsetTilePerThread = 8; + +static constexpr int NumThreadsHist = 256; +static constexpr int NumWarpsHist = NumThreadsHist / WarpSize; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +static __device__ inline float sigmoid_accurate(float x) { return 0.5f * tanhf(0.5f * x) + 0.5f; } + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +__host__ __device__ constexpr T mulLog2(T a, T bLog2) { + return a << bLog2; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +__host__ __device__ constexpr T divUpLog2(T a, T bLog2) { + return ((a + (1 << bLog2) - 1) >> bLog2); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +__host__ __device__ constexpr T divUpMulLog2(T a, T bLog2) { + return mulLog2(divUpLog2(a, bLog2), bLog2); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +__host__ __device__ constexpr int32_t getBits(int32_t value, int idx) { + int mask = idx == 0 ? 0x000000FF : idx == 1 ? 0x0000FF00 : idx == 2 ? 0x00FF0000 : 0xFF000000; + return (value & mask) >> (idx * 8); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +__host__ __device__ constexpr void setBits(int32_t& value, int32_t newBits, int idx) { + if constexpr (!IsZero) { + int mask = idx == 0 ? 0xFFFFFF00 : idx == 1 ? 0xFFFF00FF : idx == 2 ? 0xFF00FFFF : 0x00FFFFFF; + value &= mask; + } + value |= (newBits << (idx * 8)); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// +template +__device__ void initArr(int startIdx, int numElts, int stride, DataType* arr, DataType value) { + if (arr != nullptr) { + for (int i = startIdx; i < numElts; i += stride) { + arr[i] = value; + } + } +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// +template +__device__ void calcSoftmax(cg::thread_block_tile const& warp, + DataType (&scores)[VecSize]) { + DataType maxScore = DataType{-INFINITY}; + DataType sumScore = DataType{0.f}; + + // Get the max score for each token + for (int i = 0; i < VecSize; ++i) { + maxScore = scores[i] >= maxScore ? scores[i] : maxScore; + } + maxScore = cg::reduce(warp, maxScore, cg::greater()); + + // Get the summation of scores for each token +#pragma unroll + for (int i = 0; i < VecSize; ++i) { + scores[i] = static_cast(exp(scores[i] - maxScore)); + sumScore += scores[i]; + } + sumScore = cg::reduce(warp, sumScore, cg::plus()); + + // Normalize the scores +#pragma unroll + for (int i = 0; i < VecSize; ++i) { + scores[i] = static_cast(scores[i] / sumScore); + } +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +__device__ DataType calcSoftmax(cg::thread_block_tile const& warp, DataType score, + int32_t laneIdx, int32_t NumTopExperts) { + DataType maxScore = DataType{-INFINITY}; + if (laneIdx < NumTopExperts) { + maxScore = score >= maxScore ? score : maxScore; + } + maxScore = cg::reduce(warp, maxScore, cg::greater()); + + float sumScore = float{0.f}; + float newScore; + // Get the summation of scores for each token + if (laneIdx < NumTopExperts) { + newScore = static_cast(score) - static_cast(maxScore); + newScore = static_cast(exp(newScore)); + sumScore += newScore; + } + sumScore = cg::reduce(warp, sumScore, cg::plus()); + + if (laneIdx < NumTopExperts) { + score = static_cast(newScore / sumScore); + } + + return score; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// +template +__device__ void routingPermutation(KernelParams params, + PackedScoreIdx* smemPackedScoreIdx, + int32_t const warpIdx, uint32_t const clusterBlockRank) { + using OutputT = typename KernelParams::OutputT; + using TypePacked = PackedScoreIdx; + + static constexpr int MaxNumTokensSingleCluster = NumBlocksPerCluster * NumThreads; + // Number of threads in the cluster. + static constexpr int NumThreadsPerCluster = NumThreads * NumBlocksPerCluster; + // same as max num tokens + static constexpr int MaxExpandedIdxPerThread = + (MaxNumTokensSingleCluster * MaxNumTopExperts + NumThreadsPerCluster - 1) / + NumThreadsPerCluster; + + // Needed for the exclusive sum of token offsets. + // Note: the scan might include more bins than needed, with bin counts of 0 to pad + using Scan = cub::BlockScan; + __shared__ typename Scan::TempStorage tempStorage; + + uint32_t const clusterThreadIdx = NumThreads * clusterBlockRank + threadIdx.x; + auto expandedIdxSize = params.mNumTokens * params.mTopK; + + // number of experts is bounded by number of threads + __shared__ int32_t __attribute((aligned(128))) smemExpertCount[NumThreads]; + __shared__ int32_t __attribute((aligned(128))) smemExpertOffset[NumThreads]; + + // pre-fill the counts with 0 + if (threadIdx.x < params.mNumExperts) { + smemExpertCount[threadIdx.x] = 0; + } + __syncthreads(); + + // each thread keeps some number of "expanded indexes" assigned to it + // note that expanded indexes simply represent tokens here. + // for each of these, we keep the associated expert and offset within expert in registers + int32_t expertIndexes[MaxExpandedIdxPerThread]; + int32_t expertOffsets[MaxExpandedIdxPerThread]; + auto localExpertExtent = params.mNumLocalExperts << params.mLocalExpertsStrideLog2; + + // In order to avoid a serialization LDG-ATOMS-LDG-ATOMS-..., we skip multiple iterations at a + // time, and branch between a fast path without bound checks and a slow path with bound checks. + // TODO(mjoux): potentially add this back for perf tuning + // int constexpr IterStride = 4; + // static_assert(MaxExpandedIdxPerThread % IterStride == 0); + + // Define a lambda to avoid code duplication in both branches. + auto loopBody = [&](int ii, int expandedIdx) { + TypePacked scoreIdx; + if constexpr (LoadExpertIdxFromGlobal) { + scoreIdx = TypePacked{static_cast(params.mPtrExpertIdx[expandedIdx].score), + static_cast(params.mPtrExpertIdx[expandedIdx].idx)}; + } else { + TypePacked const* remoteSmem = cg::cluster_group::map_shared_rank( + smemPackedScoreIdx, expandedIdx / (NumWarps * params.mTopK)); + scoreIdx = remoteSmem[expandedIdx % (NumWarps * params.mTopK)]; + } + + expertIndexes[ii] = scoreIdx.idx; + // check whether this expert is local to our GPU at all and ignore if not + auto localExpertIdx = scoreIdx.idx - params.mLocalExpertsStartIdx; + auto isLocalExpert = localExpertIdx >= 0 && localExpertIdx < localExpertExtent && + (localExpertIdx & params.mLocalExpertsStrideLog2) == 0; + expertOffsets[ii] = isLocalExpert ? atomicAdd(smemExpertCount + scoreIdx.idx, 1) : 0; + if (params.mPtrExpertWeights != nullptr) { + params.mPtrExpertWeights[expandedIdx] = OutputT{scoreIdx.score}; + } + }; + + int constexpr IterStride = 4; +#pragma unroll + for (int32_t ii0 = 0; ii0 < MaxExpandedIdxPerThread; ii0 += IterStride) { + // Whether it's safe to do multiple iterations without bound checks. + bool const takeFastPath = (ii0 + IterStride) * NumThreadsPerCluster <= expandedIdxSize; + if (takeFastPath) { +#pragma unroll + for (int32_t jj = 0; jj < IterStride; jj++) { + int const ii = ii0 + jj; + auto expandedIdx = static_cast(clusterThreadIdx) + ii * NumThreadsPerCluster; + loopBody(ii, expandedIdx); + } + } else { + bool doBreak = false; +#pragma unroll + for (int32_t jj = 0; jj < IterStride; jj++) { + int const ii = ii0 + jj; + auto expandedIdx = static_cast(clusterThreadIdx) + ii * NumThreadsPerCluster; + if (expandedIdx >= expandedIdxSize) { + doBreak = true; + break; + } + loopBody(ii, expandedIdx); + } + if (doBreak) { + break; + } + } + } + // Make local histogram (token counts per expert) available to all threads in the cluster. + __cluster_barrier_arrive(); + __cluster_barrier_wait(); + + // + // Each thread now represents one expert + // + + // Total number of tokens for this expert. + int32_t count = 0; + // Per-expert offset for this block. + int32_t blockExpertOffset = 0; + + if (threadIdx.x < params.mNumExperts) { + // Get the histogram bin from each rank for this expert. + int32_t expertCounts[NumBlocksPerCluster]; +#pragma unroll + for (int rank = 0; rank < NumBlocksPerCluster; rank++) { + int32_t const* remoteSmem = cg::cluster_group::map_shared_rank(smemExpertCount, rank); + expertCounts[rank] = rank * NumWarps < params.mNumTokens ? remoteSmem[threadIdx.x] : 0; + } + + // Compute an exclusive prefix sum of the block-local count. +#pragma unroll + for (int rank = 0; rank < NumBlocksPerCluster; rank++) { + if (rank == clusterBlockRank) { + blockExpertOffset = count; + } + count += expertCounts[rank]; + } + } + + // Arrive: we do not access distributed shared memory after this point. + __cluster_barrier_arrive(); + + // Compute the runtime config for projections + // Whether or not an expert is local is taken into account when smemExpertCount is computed + // so we do not need to take it into account here. + const int32_t numCta = divUpLog2(count, params.mPaddingLog2); + int32_t ctaOffset; + int32_t numNonExitingCtas; + Scan(tempStorage).ExclusiveSum(numCta, ctaOffset, numNonExitingCtas); + + if (threadIdx.x < params.mNumExperts) { + // Strided loop to share this work between blocks. + for (int32_t cta = clusterBlockRank; cta < numCta; cta += NumBlocksPerCluster) { + const int32_t localExpertIdx = + (threadIdx.x - params.mLocalExpertsStartIdx) >> params.mLocalExpertsStrideLog2; + params.mPtrCtaIdxXyToBatchIdx[ctaOffset + cta] = localExpertIdx; + params.mPtrCtaIdxXyToMnLimit[ctaOffset + cta] = + min(mulLog2(ctaOffset + cta + 1, params.mPaddingLog2), + mulLog2(ctaOffset, params.mPaddingLog2) + count); + } + + // get the padded offset associated with this expert + const int32_t offset = mulLog2(ctaOffset, params.mPaddingLog2); + + // write expert offsets to shared + smemExpertOffset[threadIdx.x] = offset + blockExpertOffset; + } + + // write out padded count + if (clusterBlockRank == 0 && warpIdx == NumWarps - 1 && cute::elect_one_sync()) { + const int32_t permutedIdxSize = mulLog2(numNonExitingCtas, params.mPaddingLog2); + params.mPtrPermutedIdxSize[0] = permutedIdxSize; + params.mPtrNumNonExitingCtas[0] = numNonExitingCtas; + } + + // make expert offsets available to all threads + __syncthreads(); + + // Wait: we cannot exit while other blocks may be accessing the current block's shared memory. + // Note: I observed a perf benefit to doing this before the final loop so the compiler can + // implement break with EXIT. + __cluster_barrier_wait(); + + // trigger the secondary kernel when using PDL + // We can't do it earlier because FC1 depends on the mPtrCtaIdxXyToBatchIdx, + // mPtrCtaIdxXyToMnLimit, mPtrNumNonExitingCtas and mPtrTotalNumPaddedTokens + // TODO: this is not sufficient to ensure visibility in the next kernel! +#if !defined(PDL_PROFILE) || PDL_PROFILE == 0 + if constexpr (KernelParams::UsePdl) { + cudaTriggerProgrammaticLaunchCompletion(); + } +#endif + + // each thread has the same "expanded indexes" assigned to it as above + // at this point, we know the final offsets of experts and the offsets within + // experts, which allows writing the final index values + +#pragma unroll + for (int32_t ii = 0; ii < MaxExpandedIdxPerThread; ++ii) { + auto expandedIdx = static_cast(clusterThreadIdx) + ii * NumThreadsPerCluster; + if (expandedIdx >= expandedIdxSize) { + break; + } + auto expertIdx = expertIndexes[ii]; + // check whether this expert is local to our GPU at all + auto localExpertIdx = static_cast(expertIdx) - params.mLocalExpertsStartIdx; + auto isLocalExpert = localExpertIdx >= 0 && localExpertIdx < localExpertExtent && + (localExpertIdx & params.mLocalExpertsStrideLog2) == 0; + auto tokenIdx = expandedIdx / params.mTopK; + auto permutedIdx = + isLocalExpert ? int32_t{smemExpertOffset[expertIdx]} + expertOffsets[ii] : int32_t{-1}; + if (params.mPtrExpandedIdxToPermutedIdx != nullptr) { + params.mPtrExpandedIdxToPermutedIdx[expandedIdx] = permutedIdx; + } + if (params.mPtrPermutedIdxToTokenIdx != nullptr && isLocalExpert) { + params.mPtrPermutedIdxToTokenIdx[permutedIdx] = tokenIdx; + } + } +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// +// Two-step approach (if number of tokens exceed limits of what cluster / cooperative launch +// variants can handle): in order to minimize the amount of data to exchange through global memory, +// we will compute the local histograms in smem twice: the first kernel will get us the total number +// of tokens per expert. The second kernel will use the smem and L2 atomics to get corresponding +// element and tile offsets. +// +// Note: the histogram calculation could also be fused with routingMainKernel, but this might be +// inefficient if we have one CTA per token doing a single global atomic. +template +__global__ void __launch_bounds__(NumThreadsHist) + routingIndicesHistogramKernel(KernelParams params) { + using OutputT = typename KernelParams::OutputT; + // number of experts is bounded by number of threads + __shared__ int32_t __attribute((aligned(128))) smemExpertCount[NumThreadsHist]; + + // For unrolling. + uint32_t constexpr NumEltsPerThread = 8; + + // Pre-fill the counts with 0 + if (threadIdx.x < params.mNumExperts) { + smemExpertCount[threadIdx.x] = 0; + } + __syncthreads(); + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + // Wait on primary grid and trigger secondary kernel. + if constexpr (KernelParams::UsePdl) { + cudaGridDependencySynchronize(); + cudaTriggerProgrammaticLaunchCompletion(); + } +#endif // if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + + uint32_t const expandedIdxSize = params.mNumTokens * params.mTopK; + uint32_t const localExpertExtent = params.mNumLocalExperts << params.mLocalExpertsStrideLog2; + + uint32_t const gridBlockOffset = blockIdx.x * NumThreadsHist; + uint32_t const gridStride = gridDim.x * NumThreadsHist; + + // Define a lambda to avoid code duplication in branches. + auto loopBody = [&](int expandedIdx) { + PackedScoreIdx scoreIdx = params.mPtrExpertIdx[expandedIdx]; + // check whether this expert is local to our GPU at all and ignore if not + auto localExpertIdx = scoreIdx.idx - params.mLocalExpertsStartIdx; + auto isLocalExpert = localExpertIdx >= 0 && localExpertIdx < localExpertExtent && + (localExpertIdx & params.mLocalExpertsStrideLog2) == 0; + if (isLocalExpert) { + atomicAdd(&smemExpertCount[scoreIdx.idx], 1); + } + + if (params.mPtrExpertWeights != nullptr) { + params.mPtrExpertWeights[expandedIdx] = static_cast(scoreIdx.score); + } + }; + + // Grid-stride loop. + for (uint32_t expandedIdx0 = gridBlockOffset * NumEltsPerThread; expandedIdx0 < expandedIdxSize; + expandedIdx0 += gridStride * NumEltsPerThread) { + // Fast path if bound checks aren't necessary + if (expandedIdx0 + NumEltsPerThread * NumThreadsHist <= expandedIdxSize) { +#pragma unroll + for (uint32_t ii = 0; ii < NumEltsPerThread; ii++) { + uint32_t expandedIdx = expandedIdx0 + ii * NumThreadsHist + threadIdx.x; + loopBody(expandedIdx); + } + } else { + for (uint32_t expandedIdx = expandedIdx0 + threadIdx.x; expandedIdx < expandedIdxSize; + expandedIdx += NumThreadsHist) { + loopBody(expandedIdx); + } + } + } + __syncthreads(); + + // + // Each thread now represents one expert + // + + // Reduce histograms with atomics. + if (threadIdx.x < params.mNumExperts) { + int32_t const localExpertCount = smemExpertCount[threadIdx.x]; + atomicAdd(¶ms.mPtrExpertCounts[threadIdx.x], localExpertCount); + } +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +__global__ void __launch_bounds__(NumThreadsHist) routingIndicesOffsetsKernel(KernelParams params) { + using OutputT = typename KernelParams::OutputT; + + // number of experts is bounded by number of threads + __shared__ int32_t __attribute((aligned(128))) smemExpertOffset[NumThreadsHist]; + __shared__ int32_t __attribute((aligned(128))) smemExpertCount[NumThreadsHist]; + __shared__ int32_t __attribute((aligned(128))) smemExpertTileOffset[NumThreadsHist]; + // needed for the exclusive sum of token offsets + using Scan = cub::BlockScan; + __shared__ typename Scan::TempStorage tempStorage; + static constexpr int MaxExpandedIdxPerThread = NumEltsPerOffsetTilePerThread; + static constexpr int MaxExpandedIdxPerBlock = NumThreadsHist * MaxExpandedIdxPerThread; + + int32_t const warpIdx = __shfl_sync(0xffffffff, threadIdx.x / WarpSize, 0); + + uint32_t const expandedIdxSize = params.mNumTokens * params.mTopK; + uint32_t const numTiles = + (expandedIdxSize + MaxExpandedIdxPerBlock - 1) / (MaxExpandedIdxPerBlock); + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + // Wait on primary grid. + if constexpr (KernelParams::UsePdl) { + cudaGridDependencySynchronize(); + } +#endif // if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + + // The expert offsets are common to all tiles of all blocks. + // Load the histogram, scan it and write offsets to shared memory. + // Note: the scan is redundant in all CTAs. Would it make sense to use an intermediate kernel for + // the scan, with PDL? + + // + // Each thread represents one expert. + // + + // Get total count for this expert. + int32_t count = (threadIdx.x < params.mNumExperts) ? params.mPtrExpertCounts[threadIdx.x] : 0; + + // Compute the runtime config for projections + // Whether or not an expert is local is taken into account when the histogram is computed + // so we do not need to take it into account here. + const int32_t numCta = divUpLog2(count, params.mPaddingLog2); + int32_t ctaOffset; + int32_t numNonExitingCtas; + Scan(tempStorage).ExclusiveSum(numCta, ctaOffset, numNonExitingCtas); + + if (threadIdx.x < params.mNumExperts) { + // Get the padded offset associated with this expert + const int32_t offset = mulLog2(ctaOffset, params.mPaddingLog2); + + // Write expert offsets to shared + smemExpertOffset[threadIdx.x] = offset; + } + + // Sync to make expert offsets available to all threads. + __syncthreads(); + + // The first block writes out padded count + if (blockIdx.x == 0 && warpIdx == NumWarpsHist - 1 && cute::elect_one_sync()) { + const int32_t permutedIdxSize = mulLog2(numNonExitingCtas, params.mPaddingLog2); + params.mPtrPermutedIdxSize[0] = permutedIdxSize; + params.mPtrNumNonExitingCtas[0] = numNonExitingCtas; + } + + if (threadIdx.x < params.mNumExperts) { + // Strided loop to share this work between blocks. + for (int32_t cta = blockIdx.x; cta < numCta; cta += gridDim.x) { + const int32_t localExpertIdx = + (threadIdx.x - params.mLocalExpertsStartIdx) >> params.mLocalExpertsStrideLog2; + params.mPtrCtaIdxXyToBatchIdx[ctaOffset + cta] = localExpertIdx; + params.mPtrCtaIdxXyToMnLimit[ctaOffset + cta] = + min(mulLog2(ctaOffset + cta + 1, params.mPaddingLog2), + mulLog2(ctaOffset, params.mPaddingLog2) + count); + } + } + + // + // Now loop on indices and compute offsets. + // + + // Grid-stride loop on 1D "tiles" of input indices. + for (uint32_t tileIdx = blockIdx.x; tileIdx < numTiles; tileIdx += gridDim.x) { + if (tileIdx > 0) { + // Sync for safe reuse of smem buffers. + __syncthreads(); + } + + // Pre-fill the counts with 0 + if (threadIdx.x < params.mNumExperts) { + smemExpertCount[threadIdx.x] = 0; + } + __syncthreads(); + + // each thread keeps has some number of "expanded indexes" assigned to it + // for each of these, we keep the associated expert and offset within expert in registers + int32_t expertIndexes[MaxExpandedIdxPerThread]; + int32_t expertOffsets[MaxExpandedIdxPerThread]; + auto localExpertExtent = params.mNumLocalExperts << params.mLocalExpertsStrideLog2; + + // Define a lambda to avoid code duplication in branches. + auto loopBody = [&](int ii, int expandedIdx) { + PackedScoreIdx scoreIdx = params.mPtrExpertIdx[expandedIdx]; + expertIndexes[ii] = scoreIdx.idx; + // check whether this expert is local to our GPU at all and ignore if not + auto localExpertIdx = scoreIdx.idx - params.mLocalExpertsStartIdx; + auto isLocalExpert = localExpertIdx >= 0 && localExpertIdx < localExpertExtent && + (localExpertIdx & params.mLocalExpertsStrideLog2) == 0; + expertOffsets[ii] = isLocalExpert ? atomicAdd(smemExpertCount + scoreIdx.idx, 1) : 0; + }; + + // For all tiles but the last, all indices are in bounds. + if (tileIdx < numTiles - 1) { +#pragma unroll + for (int32_t ii = 0; ii < MaxExpandedIdxPerThread; ii += 1) { + auto expandedIdx = tileIdx * MaxExpandedIdxPerBlock + ii * NumThreadsHist + threadIdx.x; + loopBody(ii, expandedIdx); + } + } else { + // For the last tile, we need to exit the loop when out of bounds. + // In order to avoid a serialization LDG-ATOMS-LDG-ATOMS-..., we skip multiple iterations at a + // time, and branch between a fast path without bound checks and a slow path with bound checks + int constexpr IterStride = 4; + static_assert(MaxExpandedIdxPerThread % IterStride == 0); + +#pragma unroll + for (int32_t ii0 = 0; ii0 < MaxExpandedIdxPerThread; ii0 += IterStride) { + // Whether it's safe to do multiple iterations without bound checks. + bool const takeFastPath = + tileIdx * MaxExpandedIdxPerBlock + (ii0 + IterStride) * NumThreadsHist <= + expandedIdxSize; + if (takeFastPath) { +#pragma unroll + for (int32_t jj = 0; jj < IterStride; jj++) { + int const ii = ii0 + jj; + auto expandedIdx = tileIdx * MaxExpandedIdxPerBlock + ii * NumThreadsHist + threadIdx.x; + loopBody(ii, expandedIdx); + } + } else { + bool doBreak = false; +#pragma unroll + for (int32_t jj = 0; jj < IterStride; jj++) { + int const ii = ii0 + jj; + auto expandedIdx = tileIdx * MaxExpandedIdxPerBlock + ii * NumThreadsHist + threadIdx.x; + if (expandedIdx >= expandedIdxSize) { + doBreak = true; + break; + } + loopBody(ii, expandedIdx); + } + if (doBreak) { + break; + } + } + } + } + + // Make local histogram (token counts per expert) available to all threads in the block. + __syncthreads(); + + // + // Each thread now represents one expert + // + + if (threadIdx.x < params.mNumExperts) { + // Add the local bin count to the common bin count and get a per-CTA offset. We use the second + // half of the histogram buffer for this histogram, because the first half already holds the + // reduced histogram from the previous kernel. + int32_t const localExpertCount = smemExpertCount[threadIdx.x]; + int32_t const tileExpertOffset = + atomicAdd(¶ms.mPtrExpertCounts[params.mNumExperts + threadIdx.x], localExpertCount); + + // Make per-expert tile offsets available to all threads in the block. + smemExpertTileOffset[threadIdx.x] = tileExpertOffset + smemExpertOffset[threadIdx.x]; + } + __syncthreads(); + + // Add tile offset and element offset and write to global memory. + auto storeLoopBody = [&](int ii, int expandedIdx) { + int32_t expertIdx = expertIndexes[ii]; + // check whether this expert is local to our GPU at all + auto localExpertIdx = static_cast(expertIdx) - params.mLocalExpertsStartIdx; + auto isLocalExpert = localExpertIdx >= 0 && localExpertIdx < localExpertExtent && + (localExpertIdx & params.mLocalExpertsStrideLog2) == 0; + auto tokenIdx = expandedIdx / params.mTopK; + auto permutedIdx = + isLocalExpert ? (expertOffsets[ii] + smemExpertTileOffset[expertIdx]) : int32_t{-1}; + if (params.mPtrExpandedIdxToPermutedIdx != nullptr) { + params.mPtrExpandedIdxToPermutedIdx[expandedIdx] = permutedIdx; + } + if (params.mPtrPermutedIdxToTokenIdx != nullptr && isLocalExpert) { + params.mPtrPermutedIdxToTokenIdx[permutedIdx] = tokenIdx; + } + }; + // Bound checks only in last tile. + if (tileIdx < numTiles - 1) { +#pragma unroll + for (int32_t ii = 0; ii < MaxExpandedIdxPerThread; ii += 1) { + auto expandedIdx = tileIdx * MaxExpandedIdxPerBlock + ii * NumThreadsHist + threadIdx.x; + storeLoopBody(ii, expandedIdx); + } + } else { +#pragma unroll + for (int32_t ii = 0; ii < MaxExpandedIdxPerThread; ii += 1) { + auto expandedIdx = tileIdx * MaxExpandedIdxPerBlock + ii * NumThreadsHist + threadIdx.x; + if (expandedIdx >= expandedIdxSize) { + break; + } + storeLoopBody(ii, expandedIdx); + } + } + } + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) +// Trigger secondary kernel. +// Note: this does not guarantee the visibility of prior writes unless the consumer executes a +// dependency sync. +#if !defined(PDL_PROFILE) || PDL_PROFILE == 0 + if constexpr (KernelParams::UsePdl) { + cudaTriggerProgrammaticLaunchCompletion(); + } +#endif +#endif // if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) +} + +} // namespace routing +} // namespace moe::dev diff --git a/include/flashinfer/trtllm/fused_moe/RoutingKernel.h b/include/flashinfer/trtllm/fused_moe/RoutingKernel.h new file mode 100644 index 0000000000000000000000000000000000000000..3b1729922d99e8f0543f986820c46923441e0ccf --- /dev/null +++ b/include/flashinfer/trtllm/fused_moe/RoutingKernel.h @@ -0,0 +1,285 @@ +/* + * Copyright (c) 2022-2025, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include + +#include "IntFastDiv.h" +#include "flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/DtypeDecl.h" +#include "flashinfer/trtllm/common/cudaUtils.h" + +namespace moe::dev { + +namespace routing { + +namespace tg = batchedGemm::trtllm::gen; + +template +struct PackedScoreIdx { + DataType score; + int16_t idx; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +struct DataBase { + bool mUsePdl{false}; + + // optional: only used as an intermediate buffer when the number of tokens is large. + // dim: max([2*NumThreads] = [512], mNumExperts*2) + int32_t* mPtrExpertCounts{nullptr}; + // optional: if `nullptr`, it is not filled + // dim: [1] + int32_t* mPtrPermutedIdxSize{nullptr}; + // optional: if `nullptr`, it is not filled + // dim: [mNumTokens * mTopK] + int32_t* mPtrExpandedIdxToPermutedIdx{nullptr}; + // optional: if `nullptr`, it is not filled + // dim: [mNumTokens * mTopK + (mNumExperts << mPaddingLog2) - mNumExperts] + // Note: this array (mPtrPermutedIdxToTokenIdx) is uninitialized + // Any out-of-bounds values are undefined. + int32_t* mPtrPermutedIdxToTokenIdx{nullptr}; + // optional: if `nullptr`, it is not filled + // dim: [mNumTokens, mTopK] + void* mPtrExpertWeights{nullptr}; + // optional: if `nullptr`, scores are used directly as input. + // If it is given, it must represent a packed value s.t. the most significant + // 16/32 bits represent the score without sigmoid activation and + // the least significant 16 bits represent the index of the chosen expert (unsigned). + // note: this is required if the number of tokens is large. + // dim: [mNumTokens, mTopK] + void* mPtrExpertIdx{nullptr}; + // optional: if `nullptr`, `mPtrExpertIdx` must be provided. + // If it is given, it represents the scores without sigmoid activation for + // each token and expert. + // note: if it is provided, we always re-compute the top1 scores + // dim: [mNumTokens, mNumExperts] + void const* mPtrScores{nullptr}; + + // + // Grouped Gemm Launch Config Buffers + // + int32_t* mPtrCtaIdxXyToBatchIdx{nullptr}; + int32_t* mPtrCtaIdxXyToMnLimit{nullptr}; + int32_t* mPtrNumNonExitingCtas{nullptr}; + + // + // Metadata + // + int32_t mNumTokens; + int32_t mNumExperts; + int32_t mTopK; + int32_t mPaddingLog2; + + /// For expert parallelization + int32_t mLocalExpertsStartIdx; + int32_t mLocalExpertsStrideLog2; + int32_t mNumLocalExperts; +}; + +template +struct KernelParamsBase { + using InputT = InputT_; + using OutputT = OutputT_; + static constexpr bool UsePdl = UsePdl_; + + // Public pointer members + int32_t* mPtrExpertCounts = nullptr; + int32_t* mPtrPermutedIdxSize = nullptr; + int32_t* mPtrExpandedIdxToPermutedIdx = nullptr; + int32_t* mPtrPermutedIdxToTokenIdx = nullptr; + int32_t* mPtrCtaIdxXyToBatchIdx = nullptr; + int32_t* mPtrCtaIdxXyToMnLimit = nullptr; + int32_t* mPtrNumNonExitingCtas = nullptr; + OutputT* mPtrExpertWeights = nullptr; + InputT const* mPtrScores = nullptr; + + // Public scalar members + int32_t mNumTokens = 0; + int32_t mNumExperts = 0; + + int32_t mPaddingLog2 = 0; + int32_t mLocalExpertsStartIdx = 0; + int32_t mLocalExpertsStrideLog2 = 0; + int32_t mNumLocalExperts = 0; + + // Public initialization function - make it a template to accept different Data types + template + void setBaseParams(DataType const& data) { + mPtrExpertCounts = data.mPtrExpertCounts; + mPtrPermutedIdxSize = data.mPtrPermutedIdxSize; + mPtrExpandedIdxToPermutedIdx = data.mPtrExpandedIdxToPermutedIdx; + mPtrPermutedIdxToTokenIdx = data.mPtrPermutedIdxToTokenIdx; + mPtrCtaIdxXyToBatchIdx = data.mPtrCtaIdxXyToBatchIdx; + mPtrCtaIdxXyToMnLimit = data.mPtrCtaIdxXyToMnLimit; + mPtrNumNonExitingCtas = data.mPtrNumNonExitingCtas; + mPtrExpertWeights = static_cast(data.mPtrExpertWeights); + mPtrScores = (InputT const*)data.mPtrScores; + + mNumTokens = data.mNumTokens; + mNumExperts = data.mNumExperts; + + mPaddingLog2 = data.mPaddingLog2; + mLocalExpertsStartIdx = data.mLocalExpertsStartIdx; + mLocalExpertsStrideLog2 = data.mLocalExpertsStrideLog2; + mNumLocalExperts = data.mNumLocalExperts; + } +}; + +namespace routingDeepSeek { + +//////////////////////////////////////////////////////////////////////////////////////////////////// +struct Data : public DataBase { + tg::Dtype mDtypeExpW{tg::Dtype::Bfloat16}; + + // + // Grouped Gemm Launch Config Buffers + // + void const* mPtrRoutingBias; + + int32_t mHiddenDim; // not used + int32_t mNumExpertGroups; + int32_t mNumLimitedGroups; + + float mRouteScale; + bool mUseRoutingSoftmax; +}; + +template +struct KernelParams : public KernelParamsBase { + using InputT = InputT_; + using OutputT = OutputT_; + + static constexpr bool UseGroups = UseGroups_; + + PackedScoreIdx* mPtrExpertIdx = nullptr; + + // OutputT* mPtrExpertWeightsFull = nullptr; + // Note: this variable(mPtrExpertWeightsFull) might need to be added back for the low-latency + // kernels for MoE in tllm-gen in the future + + OutputT const* mPtrRoutingBias = nullptr; + + int32_t mNumExpertGroups = 0; + int32_t mNumExpertsPerGroup = 0; + int32_t mNumLimitedGroups = 0; + + trtllm::dev::IntFastDiv mTopK; + float mRouteScale = 0.f; + + static KernelParams setKernelParams(Data const& data) { + KernelParams params; + params.setBaseParams(data); + + params.mPtrExpertIdx = (PackedScoreIdx*)data.mPtrExpertIdx; + + // params.mPtrExpertWeightsFull = static_cast(data.mPtrExpertWeightsFull); + params.mPtrRoutingBias = static_cast(data.mPtrRoutingBias); + + params.mNumExpertGroups = data.mNumExpertGroups; + params.mNumExpertsPerGroup = data.mNumExperts / data.mNumExpertGroups; + params.mNumLimitedGroups = data.mNumLimitedGroups; + params.mTopK = trtllm::dev::IntFastDiv(data.mTopK); + params.mRouteScale = data.mRouteScale; + + return params; + } +}; + +void run(Data& data, void* stream); + +} // namespace routingDeepSeek + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace routingLlama4 { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +struct Data : public DataBase { + tg::Dtype mDtypeExpW{tg::Dtype::Bfloat16}; +}; + +template +struct KernelParams : public KernelParamsBase { + using InputT = InputT_; + using OutputT = OutputT_; + + PackedScoreIdx* mPtrExpertIdx = nullptr; + + int32_t mTopK; + + static KernelParams setKernelParams(Data const& data) { + KernelParams params; + params.setBaseParams(data); + + params.mPtrExpertIdx = (PackedScoreIdx*)data.mPtrExpertIdx; + params.mTopK = data.mTopK; + return params; + } +}; + +void run(Data const& data, void* stream); + +} // namespace routingLlama4 + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace routingRenormalize { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +struct Data : public DataBase { + tg::Dtype mDtypeExpW{tg::Dtype::Fp32}; + tg::Dtype mDtypeElt{tg::Dtype::Bfloat16}; + + bool mDoSoftmaxBeforeTopK{false}; + bool mNormTopkProb{true}; // Default value is true for Qwen3 model +}; + +template +struct KernelParams : public KernelParamsBase { + using InputT = InputT_; + using OutputT = OutputT_; + + static constexpr bool DoSoftmaxBeforeTopK = DoSoftmaxBeforeTopK_; + + PackedScoreIdx* mPtrExpertIdx = nullptr; + + int32_t mTopK = 0; + + bool mNormTopkProb = true; + + static KernelParams setKernelParams(Data const& data) { + KernelParams params; + params.setBaseParams(data); + + params.mPtrExpertIdx = (PackedScoreIdx*)data.mPtrExpertIdx; + params.mNormTopkProb = data.mNormTopkProb; + params.mTopK = data.mTopK; + return params; + } +}; + +void run(Data const& data, void* stream); + +} // namespace routingRenormalize + +//////////////////////////////////////////////////////////////////////////////////////////////////// +} // namespace routing +} // namespace moe::dev diff --git a/include/flashinfer/trtllm/fused_moe/RoutingKernelTopK.cuh b/include/flashinfer/trtllm/fused_moe/RoutingKernelTopK.cuh new file mode 100644 index 0000000000000000000000000000000000000000..e3ea61eb7b412f5128f61b7dcb6730603807790f --- /dev/null +++ b/include/flashinfer/trtllm/fused_moe/RoutingKernelTopK.cuh @@ -0,0 +1,196 @@ +/* + * Copyright (c) 2022-2025, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include + +#include + +namespace moe::dev::routing { + +namespace topk { + +//////////////////////////////////////////////////////////////////////////////////////////////////// +namespace cg = cooperative_groups; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +static constexpr int WarpSize = 32; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct TopKRedType { + using TypeExpW = TypeExpW_; + static_assert(std::is_same_v || std::is_same_v || + std::is_same_v, + "Top K reduction only implemented for float, float16 and bfloat16"); + + using TypeCmp = std::conditional_t; + using IdxT = std::conditional_t; + static constexpr int moveBits = (sizeof(TypeExpW) == 4) ? 32 : 16; + static constexpr int maxIdx = 65535; + + TypeCmp compVal; + + static __host__ __device__ inline TypeCmp makeCmpVal(TypeExpW val, int32_t idx = 0) { + auto valueBits = cub::Traits::TwiddleIn( + reinterpret_cast::UnsignedBits&>(val)); + TypeCmp compactTmp = reinterpret_cast(valueBits); + compactTmp = (compactTmp << moveBits) | (0xFFFF & (maxIdx - idx)); + // Use 65535 minus idx to give higher priority to elements with smaller indices. + return compactTmp; + } + + static __host__ __device__ inline void unpack(TypeExpW& value, int32_t& index, TypeCmp cmp) { + // Since idx is always smaller than 65536 and positive, we can directly use it as the lower 16 + // bits + index = maxIdx - static_cast(cmp & 0xFFFF); + + auto compactTmp = cmp >> moveBits; + auto valueBits = cub::Traits::TwiddleOut( + reinterpret_cast::UnsignedBits&>(compactTmp)); + value = reinterpret_cast(valueBits); + } + + __host__ __device__ TopKRedType() = default; + + __host__ __device__ TopKRedType(TypeExpW val, int32_t idx) : compVal(makeCmpVal(val, idx)) {} + + __host__ __device__ operator TypeCmp() const noexcept { return compVal; } + + __device__ inline TypeCmp reduce(cg::thread_block_tile const& warp) { +#ifdef __CUDA_ARCH__ + static constexpr bool hasFastRedux = __CUDA_ARCH__ >= 1000; +#else + static constexpr bool hasFastRedux = false; +#endif + if constexpr (!hasFastRedux || sizeof(TypeCmp) == 8) { + return cg::reduce(warp, compVal, cg::greater{}); + } else { + TypeCmp result; + asm("redux.sync.max.u32 %0, %1, 0xffffffff;\n" : "=r"(result) : "r"(compVal)); + return result; + } + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +#define TOPK_SWAP(I, J) \ + { \ + auto pairMin = min(topK[I].compVal, topK[J].compVal); \ + auto pairMax = max(topK[I].compVal, topK[J].compVal); \ + topK[I].compVal = pairMax; \ + topK[J].compVal = pairMin; \ + } + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct Sort; + +template +struct Sort<1, RedType> { + static __device__ void run(RedType* topK) {} +}; + +template +struct Sort<2, RedType> { + static __device__ void run(RedType* topK) { TOPK_SWAP(0, 1); } +}; + +template +struct Sort<3, RedType> { + static __device__ void run(RedType* topK) { + TOPK_SWAP(0, 1); + TOPK_SWAP(1, 2); + TOPK_SWAP(0, 1); + } +}; + +template +struct Sort<4, RedType> { + static __device__ void run(RedType* topK) { + TOPK_SWAP(0, 2); + TOPK_SWAP(1, 3); + TOPK_SWAP(0, 1); + TOPK_SWAP(2, 3); + TOPK_SWAP(1, 2); + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +__forceinline__ __device__ void reduceTopK(cg::thread_block_tile const& warp, + Type (&out)[K], int32_t (&outIdx)[K], Type value, + int32_t idx, Type const minValue, int actualK = K) { + static_assert(K > 0, "Top K must have K > 0"); + static_assert(K < WarpSize, "Top K must have K < WarpSize"); + using RedType = TopKRedType; + RedType topK{value, idx}; + typename RedType::TypeCmp packedMax{}; +#pragma unroll + for (int kk = 0; kk < actualK; ++kk) //@todo: check if actualK is correct + { + topK = kk > 0 && packedMax == topK.compVal ? RedType{minValue, idx} : topK; + // get the next largest value + packedMax = topK.reduce(warp); + RedType::unpack(out[kk], outIdx[kk], packedMax); + } +}; + +template +__forceinline__ __device__ void reduceTopK(cg::thread_block_tile const& warp, + Type (&out)[K], int32_t (&outIdx)[K], Type (&value)[N], + int32_t (&idx)[N], Type const minValue, + int actualK = K) { + static_assert(K > 0, "Top K must have K > 0"); + static_assert(K < WarpSize, "Top K must have K < WarpSize"); + static_assert(N > 0, "Top K must have N > 0"); + static_assert(N < 5, "Only support candidates number less than or equal to 128"); + using RedType = TopKRedType; + RedType topK[N]; +#pragma unroll + for (int nn = 0; nn < N; ++nn) { + topK[nn] = RedType{value[nn], idx[nn]}; + } + + Sort::run(topK); + + typename RedType::TypeCmp packedMax{}; +#pragma unroll + for (int kk = 0; kk < actualK; ++kk) //@todo: check if actualK is correct + { + bool update = kk > 0 && packedMax == topK[0].compVal; +#pragma unroll + for (int nn = 0; nn < N; ++nn) { + topK[nn] = update && nn == N - 1 ? RedType{minValue, idx[nn]} + : update ? topK[nn + 1] + : topK[nn]; + } + // get the next largest value + packedMax = topK[0].reduce(warp); + RedType::unpack(out[kk], outIdx[kk], packedMax); + } +}; + +#undef TOPK_SWAP +} // namespace topk +} // namespace moe::dev::routing diff --git a/include/flashinfer/trtllm/fused_moe/runner.h b/include/flashinfer/trtllm/fused_moe/runner.h new file mode 100644 index 0000000000000000000000000000000000000000..972b8d0bba5456eaa231444e466d2430e3c4b13a --- /dev/null +++ b/include/flashinfer/trtllm/fused_moe/runner.h @@ -0,0 +1,356 @@ +/* + * Copyright (c) 2022-2025, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +#include "DevKernel.h" +#include "RoutingKernel.h" +// #include "flashinfer/trtllm/common/cudaDriverWrapper.h" +#include "flashinfer/trtllm/batched_gemm/KernelRunner.h" +#include "flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/DtypeDecl.h" +#include "flashinfer/trtllm/common/cudaUtils.h" + +namespace tensorrt_llm { +namespace kernels { +namespace trtllmgen_moe { + +namespace Routing { + +// The type of method in top-K routing, for use in torch custom op +// Please keep this in sync with the counterpart defined in +// tensorrt_llm/_torch/modules/fused_moe/routing.py +enum class RoutingMethodType : int64_t { + // Default: Softmax -> TopK + Default = 0, + // Renormalize: TopK -> Softmax + Renormalize = 1, + // DeepSeekV3: Sigmoid -> RoutingBiasAdd -> Top2 in group -> Top4 groups -> Top8 experts from the + // Top4 groups + DeepSeekV3 = 2, + // Llama4: Top1 -> Sigmoid + Llama4 = 3, + // RenormalizeNaive: Softmax -> TopK -> Renormalize + RenormalizeNaive = 4, + // Unspecified + Unspecified = 5, +}; + +inline std::string serializeMoeRoutingMethodType(RoutingMethodType routingMethodType) { + switch (routingMethodType) { + case RoutingMethodType::Default: + return "Default"; + case RoutingMethodType::Renormalize: + return "Renormalize"; + case RoutingMethodType::DeepSeekV3: + return "DeepSeekV3"; + case RoutingMethodType::Llama4: + return "Llama4"; + case RoutingMethodType::RenormalizeNaive: + return "RenormalizeNaive"; + default: + return "InvalidRountingMethod"; // TODO throw error + }; +} + +inline int32_t getMaxPermutedPaddedCount(int32_t numTokens, int32_t expertsPerToken, + int32_t numExperts, int32_t padding) { + auto const expandedRowCount = numTokens * expertsPerToken; + auto const maxPaddingRequired = (padding - 1) * numExperts; + return common::roundUp(expandedRowCount + maxPaddingRequired, padding); +} + +inline int32_t getMaxNumCtasInBatchDim(int32_t numTokens, int32_t topK, int32_t numExperts, + int32_t tileTokensDim) { + // Get maximum number of CTAs in batch dim per expert. + auto const maxCtasInBatchDimPerExpert = common::ceilDiv(numTokens, tileTokensDim); + // Get maximum enabled experts. + auto const maxEnabledExperts = std::min(numTokens * topK, numExperts); + // Get maximum number of CTAs in batch dim. + auto maxNumCtasInBatchDim = maxEnabledExperts * maxCtasInBatchDimPerExpert; + + // For large token counts, the above bound can be pessimistic since not all the tokens can + // be routed to all the enabled experts. Instead we can essentially bound the number of CTAs + // by permuted buffer size. However, this method will be overly pessimistic for low-token + // counts + auto const tilesForPermutedBuffer = common::ceilDiv( + getMaxPermutedPaddedCount(numTokens, topK, numExperts, tileTokensDim), tileTokensDim); + + // Set maxNumCtasInBatchDim to be the minimum of the two methods + maxNumCtasInBatchDim = std::min(maxNumCtasInBatchDim, tilesForPermutedBuffer); + + return maxNumCtasInBatchDim; +} + +class Runner { + public: + explicit Runner(); + + explicit Runner(int32_t tileTokensDim); + + void run(void* routingLogits, void* routingBias, int32_t numTokens, int32_t numExperts, + int32_t topK, int32_t nGroups, int32_t topkGroups, int32_t localExpertOffset, + int32_t localNumExperts, float routedScalingFactor, int32_t* routingExpertIndexes, + int32_t* expertCountHistogram, int32_t* permutedIdxSize, + int32_t* expandedIdxToPermutedIdx, int32_t* permutedIdxToExpandedIdx, + int32_t* permutedIdxToTokenIdx, void* expertWeights, int32_t* numTokensPerExpert, + int32_t* ctaIdxXyToBatchIdx, int32_t* ctaIdxXyToMnLimit, int32_t* numNonExitingCtas, + batchedGemm::trtllm::gen::Dtype dtypeElt, bool useRoutingScalesOnInput, + bool useDeepSeekFp8, RoutingMethodType routingMethodType, cudaStream_t stream); + + private: + int32_t mTileTokensDim{8}; +}; +} // namespace Routing + +namespace PermuteGemm1 { +class Runner { + public: + explicit Runner(batchedGemm::trtllm::gen::Dtype dtypeAct, + batchedGemm::trtllm::gen::Dtype dtypeWeights, bool useDeepSeekFp8, + int tileTokensDim, ActType actType, bool useShuffledMatrixA, + batchedGemm::gemm::MatrixLayout weight_layout); + + size_t getWorkspaceSizeInBytes(int32_t topK, int32_t hiddenSize, int32_t intermediateSize, + int32_t numExperts, int32_t numTokens, int32_t configIndex) const; + + [[nodiscard]] int32_t getDefaultValidConfigIndex(int32_t topK, int32_t hiddenSize, + int32_t intermediateSize, int32_t numExperts, + int32_t numTokens) const; + + [[nodiscard]] bool isValidConfigIndex(int32_t configIndex, int32_t topK, int32_t hiddenSize, + int32_t intermediateSize, int32_t numExperts, + int32_t numTokens) const; + + [[nodiscard]] std::vector getPassingConfigIndices() const; + + void run(void* hiddenState, void* hiddenStateScale, void* weight, void* weightScale, + void* expertWeights, float* outputScalesScalar, float* outputScalesGateScalar, + float* ptrBias, float* ptrSwiGluAlpha, float* ptrSwiGluBeta, float* ptrClampLimit, + void* output, void* outputScale, int32_t topK, int32_t hiddenSize, + int32_t intermediateSize, int32_t numExperts, int32_t numTokens, + int32_t* permutedIdxToTokenIdx, int32_t* ptrNumNonExitingCtas, + int32_t* ptrTotalNumPaddedTokens, int32_t* ptrCtaIdxXyToBatchIdx, + int32_t* ptrCtaIdxXyToMnLimit, void* bmm1Workspace, bool useRoutingScalesOnInput, + int device, cudaStream_t stream, int32_t configIndex); + + private: + batchedGemm::trtllm::gen::Dtype mDtypeAct; + batchedGemm::trtllm::gen::Dtype mDtypeWeights; + int32_t mTileTokensDim; + tensorrt_llm::kernels::TrtllmGenBatchedGemmRunner mRunner; +}; +} // namespace PermuteGemm1 + +namespace Gemm2 { +class Runner { + public: + explicit Runner(batchedGemm::trtllm::gen::Dtype dtypeAct, + batchedGemm::trtllm::gen::Dtype dtypeWeights, + batchedGemm::trtllm::gen::Dtype outputDtype, bool useDeepSeekFp8, + int tileTokensDim, bool useShuffledMatrixA, + batchedGemm::gemm::MatrixLayout weight_layout); + + size_t getWorkspaceSizeInBytes(int32_t topK, int32_t hiddenSize, int32_t intermediateSize, + int32_t numExperts, int32_t numTokens, int32_t configIndex) const; + + [[nodiscard]] int32_t getDefaultValidConfigIndex(int32_t topK, int32_t hiddenSize, + int32_t intermediateSize, int32_t numExperts, + int32_t numTokens) const; + + [[nodiscard]] bool isValidConfigIndex(int32_t configIndex, int32_t topK, int32_t hiddenSize, + int32_t intermediateSize, int32_t numExperts, + int32_t numTokens) const; + + [[nodiscard]] std::vector getPassingConfigIndices() const; + + void run(void* permutedHiddenState, void* permutedHiddenStateScale, void* weight, + void* weightScale, float* outputScalesScalar, float* ptrBias, void* output, + void* outputScale, int32_t topK, int32_t hiddenSize, int32_t intermediateSize, + int32_t numExperts, int32_t numTokens, int32_t* ptrNumNonExitingCtas, + int32_t* ptrTotalNumPaddedTokens, int32_t* ptrCtaIdxXyToBatchIdx, + int32_t* ptrCtaIdxXyToMnLimit, void* bmm2Workspace, int device, cudaStream_t stream, + int32_t configIndex); + + private: + batchedGemm::trtllm::gen::Dtype mDtypeAct; + batchedGemm::trtllm::gen::Dtype mDtypeWeights; + batchedGemm::trtllm::gen::Dtype mDtypeOut; + int32_t mTileTokensDim; + tensorrt_llm::kernels::TrtllmGenBatchedGemmRunner mRunner; +}; +} // namespace Gemm2 + +namespace MoE { +namespace btg = batchedGemm::trtllm::gen; + +struct MoERunnerArgs { + void* routing_logits = nullptr; // [num_tokens, num_experts] in float, generated after + // gemm(hidden_state, routing_weights) + void* routing_bias = nullptr; // [num_experts] in bfloat16 for now = mDtypeExpW + void* hidden_states = nullptr; // [num_tokens, hidden_size] in fp8 = mDtypeElt + // [hidden_size/128, num_tokens] in float for e4m3 DS recipe + // and [num_tokens, hidden_size/16] in float for e2m1 + void* hidden_states_scale = nullptr; + + // Gemm input: + void* gemm1_weights = nullptr; + void* gemm1_weights_scale = nullptr; + void* gemm2_weights = nullptr; + void* gemm2_weights_scale = nullptr; + + float* gemm1_bias = nullptr; + float* gemm1_alpha = nullptr; + float* gemm1_beta = nullptr; + float* gemm1_clamp_limit = nullptr; + float* gemm2_bias = nullptr; + + int32_t num_tokens{0}; + int32_t num_experts{0}; + // Hidden dimension input of MoE block. It might be padded. + int32_t hidden_size{0}; + // Hidden dimension output of MoE block. It is not padded. + // If not provided it is the same as hidden_size. + std::optional hidden_size_output; + // TODO: only compiled routing kernel supports top_k = 8 + int32_t top_k{0}; + int32_t n_group{0}; + // TODO: only compiled routing kernel supports topk_group = 4 + int32_t topk_group{0}; + float routed_scaling_factor{0.0f}; + int32_t intermediate_size{0}; + int32_t local_expert_offset{0}; + int32_t local_num_experts{0}; + // TODO: support other types + btg::Dtype mDtypeElt{btg::Dtype::Void}; + btg::Dtype mDtypeExpW{btg::Dtype::Bfloat16}; + btg::Dtype mDtypeOut{btg::Dtype::Bfloat16}; + + // Apply routing scale factors to input activations + bool mUseRoutingScalesOnInput{false}; + bool mUseDeepSeekFp8{false}; + float* output1_scales_scalar = nullptr; + float* output1_scales_gate_scalar = nullptr; + float* output2_scales_scalar = nullptr; + + // Output: + void* output = nullptr; + float* output_scale = nullptr; + + // finalize + bool do_finalize{true}; +}; + +struct MoEWorkspace { + // Routing intermediate outputs: + int32_t* routing_expert_indexes = nullptr; + int32_t* permuted_idx_size = nullptr; + int32_t* total_num_padded_tokens = nullptr; // TODO: duplicate of permuted_idx_size + int32_t total_max_padded_tokens{0}; + + int32_t* expanded_idx_to_permuted_idx = nullptr; + int32_t* permuted_idx_to_expanded_idx = nullptr; + int32_t* permuted_idx_to_token_idx = nullptr; + void* expert_weights = nullptr; // [num_tokens, top_k] in bfloat16 = mDtypeExpW + + int32_t* cta_idx_xy_to_batch_idx = nullptr; + int32_t* cta_idx_xy_to_mn_limit = nullptr; + int32_t* num_non_exiting_ctas = nullptr; + + void* hidden_states_scale_linear = nullptr; + + // Permute intermediate outputs: + void* permuted_hidden_states = nullptr; + float* permuted_hidden_states_scale = nullptr; + + // Gemm1 intermediate outputs: + int32_t ProjUpTileN{0}; + void* gemm1_output = nullptr; + float* gemm1_output_scale = nullptr; + + // Activation intermediate outputs: + void* activation_output = nullptr; + float* activation_output_scale = nullptr; + + // Gemm2 intermediate outputs: + void* gemm2_output = nullptr; + float* gemm2_output_scale = nullptr; + + // Finalize intermediate outputs (placeholder not used) + void* finalize_output = nullptr; + float* finalize_output_scale = nullptr; + + // FC1 workspace: + void* bmm1_workspace = nullptr; + + // FC2 workspace: + void* bmm2_workspace = nullptr; +}; + +// Config indices to be used with Batched GEMM runners +struct MoEConfig { + int64_t gemm1Config; + int64_t gemm2Config; +}; + +class Runner { + public: + // FIXME: tileTokensDim is hardcoded for now + Runner(batchedGemm::trtllm::gen::Dtype dtypeAct, batchedGemm::trtllm::gen::Dtype dtypeWeights, + bool useDeepSeekFp8, int tileTokensDim = 8, ActType actType = ActType::SwiGlu, + bool useShuffledMatrixA = false, + batchedGemm::gemm::MatrixLayout weight_layout = batchedGemm::gemm::MatrixLayout::MajorK); + Runner(batchedGemm::trtllm::gen::Dtype dtypeElt, bool useDeepSeekFp8, int tileTokensDim = 8, + bool useShuffledMatrixA = false, + batchedGemm::gemm::MatrixLayout weight_layout = batchedGemm::gemm::MatrixLayout::MajorK); + + void run(MoERunnerArgs const& args, MoEWorkspace const& workspace, int device, + cudaStream_t stream, int64_t configIndex); + + [[nodiscard]] std::tuple getWorkspaceSizeInBytes(MoERunnerArgs const& args, + int64_t configIndex) const; + + [[nodiscard]] std::vector getValidConfigIndices(int32_t topK, int32_t hiddenSize, + int32_t intermediateSize, + int32_t numLocalExperts, + int32_t numTokens) const; + + [[nodiscard]] int64_t getDefaultValidConfigIndex(int32_t topK, int32_t hiddenSize, + int32_t intermediateSize, + int32_t numLocalExperts, + int32_t numTokens) const; + + private: + void setOpsData(MoERunnerArgs const& args, MoEWorkspace const& workspace, + moe::dev::convertsf::Data& convertSfData, + moe::dev::activation::Data& activationData, + moe::dev::finalize::Data& finalizeData); + + private: + PermuteGemm1::Runner mPermuteGemm1; + Gemm2::Runner mGemm2; + + // This will be the cartesian product of the passing configs for gemm1 and gemm2 + // This allows us to autotune the MoE as one operation instead of tuning gemm1 and gemm2 + // separately + std::vector mPassingConfigs; +}; +} // namespace MoE + +} // namespace trtllmgen_moe +} // namespace kernels +} // namespace tensorrt_llm diff --git a/include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/Enums.h b/include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/Enums.h new file mode 100644 index 0000000000000000000000000000000000000000..8439d82de358e52510dfeffa83f38a6035bac2d8 --- /dev/null +++ b/include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/Enums.h @@ -0,0 +1,128 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & + * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include + +namespace gemm { + +namespace gemm { + +enum class AllReduceAlgo : uint32_t { + // Does not apply all-reduce. + None = 0, + // Reduction occurs at L2 cache; pulls N-1 partial outputs from peer devices. Result is + // non-deterministic. Potentially lower latency at cost of higher memory traffic. + OneShot, + // Reduction occurs at switch; pulls 1/Nth of the output from switch (reduce-scatter phase) and + // store to multicast mem (all-gather phase). Result is deterministic. Lower memory traffic at + // cost of potentially higher latency. + TwoShot, +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +enum class MatrixLayout { + // K-major layout (default). [Mn, K] + MajorK = 0, + // M-major for A and N-major for B. [K, Mn] + MajorMn, + // Layout is blocked along the K dimension as seen in the diagram below. [K / blockK, Mn, blockK] + // where blockK is fixed at 128B + // + // ├────────────── K ──────────────┤ + // ┬ ┬ ├──── K block ───┤ + // │ │ │ 0 1 2 3 â•‘ 32 33 34 35 │ + // │ CTA0 │ 4 5 6 7 â•‘ 36 37 38 39 │ + // │ │ │ 8 9 10 11 â•‘ 40 41 42 43 │ + // │ â”´ │ 12 13 14 15 â•‘ 44 45 46 47 │ + // M ┬ ├────────────────║────────────────┤ + // │ │ │ 16 17 18 19 â•‘ 48 49 50 51 │ + // │ CTA1 │ 20 21 22 23 â•‘ 52 53 54 55 │ + // │ │ │ 24 25 26 27 â•‘ 56 57 58 59 │ + // â”´ â”´ │ 28 29 30 31 â•‘ 60 61 62 63 │ + BlockMajorK +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +enum class SplitK : uint32_t { + // No split-k is needed. I.e. mNumSlicesForSplitK == 1. + None = 0, + // CTAs computing one MN tile save partial results to global memory. + // Then wait on the barrier and the last CTA in the group loads partial results from gmem, + // sums them up and writes back to gmem. + Gmem, + // All CTAs in one CGA calculate partial sums. Then send the results to the smem of + // the last CTA in the CGA, which sums them up and writes to gmem. + Dsmem, +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +enum class BiasType : uint32_t { + // No bias. + None = 0, + // One bias value per N of the output tensor. + M = 1, + // One bias value per row M of the output tensor. + N = 2, + // One bias value for each element of the output tensor. + Mn = 3, +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +enum class TileScheduler { + // Static scheduler (Non-persistent). + Static = 0, + // Dynamic persistent scheduler. This is either based on an atomically incremented global work id + // prior to SM100 archs, or the HW supported work id scheduler based on UGETNEXTWORKID for SM100+. + Persistent, +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// Helper functions to check the SplitK type. + +#define SPLIT_K_FUNCTION(Mode) \ + inline bool doesSplitKUse##Mode(SplitK mode) { return (mode == SplitK::Mode); } + +SPLIT_K_FUNCTION(Gmem) +SPLIT_K_FUNCTION(Dsmem) + +#undef SPLIT_K_FUNCTION + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// Helper functions to check the Bias type. + +#define BIAS_TYPE_FUNCTION(Mode) \ + inline bool isBiasType##Mode(BiasType type) { return (type == BiasType::Mode); } + +BIAS_TYPE_FUNCTION(None) +BIAS_TYPE_FUNCTION(N) +BIAS_TYPE_FUNCTION(M) +BIAS_TYPE_FUNCTION(Mn) + +#undef BIAS_TYPE_FUNCTION + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace gemm + +} // namespace gemm diff --git a/include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/GemmInterface.h b/include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/GemmInterface.h new file mode 100644 index 0000000000000000000000000000000000000000..ba5d7c06db135cea991f8a1f902f3df531f24a18 --- /dev/null +++ b/include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/GemmInterface.h @@ -0,0 +1,540 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & + * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include +#include + +#include "GemmOptions.h" +#include "KernelParams.h" +#include "trtllm/gen/CudaKernelLauncher.h" + +#ifdef TLLM_GEN_EXPORT_INTERFACE +#include "KernelMetaInfo.h" +#endif // TLLM_GEN_EXPORT_INTERFACE + +#ifdef TLLM_GEN_GEMM_CUBIN_PATH +static const std::string tllm_gen_gemm_cubin_path = std::string(TLLM_GEN_GEMM_CUBIN_PATH); +#else +static_assert(false, "TLLM_GEN_GEMM_CUBIN_PATH macro is not defined when compiling"); +#endif + +namespace flashinfer::trtllm_cubin_loader { +std::string getCubin(const std::string& kernelName, const std::string& sha256); +} // namespace flashinfer::trtllm_cubin_loader + +namespace gemm { + +namespace gemm { + +//////////////////////////////////////////////////////////////////////////////////////////////////// +// +// GemmData +// +//////////////////////////////////////////////////////////////////////////////////////////////////// + +struct GemmData { + struct ProblemDimensions { + // The M dimension. + // It is the total number of tokens if A is the activation matrix. + // It is the total number of output channels if A is the weight matrix. + int32_t mM{0}; + // The N dimension. + // It is the total number of tokens if B is the activation matrix. + // It is the total number of output channels if B is the weight matrix. + int32_t mN{0}; + // The K dimension. It is the hidden dimension of the input matrices. + int32_t mK{0}; + // The rank id of the current device in the multi-gpu space. + int32_t mRank{0}; + // The number of devices in tensor-parallel group. + int32_t mWorldSize{0}; + }; + + struct InputBuffers { + // The matrix A. The data type is controlled by options.mDtypeA. + // + // When layoutA is MatrixLayout::MajorK, the shape is [M, K]. + // When LayoutA is MatrixLayout::MajorMn, the shape is [K, M]. + // When LayoutA is MatrixLayout::BlockMajorK, the shape is [K / blockK, M, blockK] where blockK + // is 128B. + // The rightmost dimension is contiguous in memory. + void const* mPtrA{nullptr}; + + // The block scaling factors to dequantize A. + // + // If DeepSeek FP8 recipe is used: + // If transposeMmaOutput is false, shape is [K / 128, M]. + // Otherwise, shape is [M / 128, K / 128]. + // The rightmost dimension is contiguous in memory. + // + // If DeepSeek FP8 recipe is not used, but for MxFp{4,8} and NvFp4 formats: + // The layout of scaling factors for A is always R128c4 + // M must be a multiple of 128. + // K must be a multiple of 64. + // The "logical" shape is: [M, K / 16]. + // The R128c4 layout is: [M / 128, K / 16 / 4, 512]. + // The shape we use for TMA is: [M / 128, K / 16 / 4, 2, 256]. + // Dtype is Dtype::E4m3. + // + // Otherwise should be set to nullptr. + void const* mPtrSfA{nullptr}; + + // The per-token scaling factors from scale A. + // + // This is used for either: + // * Per-token scaling factor quantization schemes, such as MetaFP8. The dtype is + // Dtype::Float32 + // * When the routing scales are applied to the input activations (only when output is not + // transposed). The dtype is Dtype::Bfloat16 + // + // The shape is [M] + void const* mPtrPerTokenSfA{nullptr}; + + // The matrix B. The data type is controlled by options.mDtypeB. + // + // When layoutB is MatrixLayout::MajorK, the shape is [N, K]. + // When layoutB is MatrixLayout::MajorMn, the shape is [K, N]. + // When layoutB is MatrixLayout::BlockMajorK, the shape is [K / blockK, N, blockK] where blockK + // is 128B. + // The rightmost dimension is contiguous in memory. + void const* mPtrB{nullptr}; + + // The scaling factors to dequantize B. + // + // If DeepSeek FP8 recipe is used: + // If transposeMmaOutput is false, shape is [N / 128, K / 128]. + // Otherwise, shape is [K / 128, N]. + // The rightmost dimension is contiguous in memory. + // + // If DeepSeek FP8 recipe is not used, but for MxFp{4,8} and NvFp4 formats: + // If the layout is R128c4, + // N must be a multiple of 128. + // K must be a multiple of 64. + // The R128c4 layout is: [N / 128, K / 16 / 4, 512] + // The shape we use for TMA is: [N / 128, K / 16 / 4, 2, 256] + // + // If the layout is R8c4, + // N must be a multiple of 8. + // K must be a multiple of 64. + // The R8c4 layout is: [N / 8, K / 16 / 4, 32] + // The shape we use for TMA is: [N / 8, K / 16 / 4 / repeats, repeats * 32] + // where repeats = min(tileK / 16 / 4, 8) + // + // Dtype is Dtype::E4m3. + // + // Otherwise should be set to nullptr. + void const* mPtrSfB{nullptr}; + + // The per-token scaling factors from scale B. + // + // This is used for either: + // * Per-token scaling factor quantization schemes, such as MetaFP8. The dtype is + // Dtype::Float32 + // * When the routing scales are applied to the input activations (only when output is + // transposed). The dtype is Dtype::Bfloat16 + // + // The shape is [N] + void const* mPtrPerTokenSfB{nullptr}; + + // The bias applied after the GEMM. + // The bias is applied before applying the global scaling factor. I.e. + // C' = (A * B + bias') * scaleC + // scaleC = dequantA * dequantB * quantC + // Thus, the bias' = bias / (dequantA * dequantB), where the bias is the original bias. + // + // if BiasType is N, the shape is [N]. + // The bias is broadcasted along the M dimension. + // + // if BiasType is M, the shape is [M]. + // The bias is broadcasted along the N dimension. + // + // The dtype is float32. + void const* mPtrBias{nullptr}; + + // The output tensor scaling factor for MxFp{4,8}, Fp8, NvFp4 and DeepSeek FP8 quantization. + // TensorRT-LLM API requires a scaling factor on the device. + // Shape is [1]. + void* mPtrScaleC{nullptr}; + }; + + struct OutputBuffers { + // The output matrix C. The data type is controlled by options.mDtypeC. + // + // When transposeMmaOutput is true, the shape is [N, M]. + // Otherwise, the shape is [M, N]. + // Elements in a given row are stored contiguously in memory (row-major). + void* mPtrC{nullptr}; + + // Pointer for output with multicast mapping. It is used by the "reduce" op (LDGMC.ADD) of the + // two-shot reduce-scatter phase. Otherwise, it should be set to nullptr. + // The shape is [M, N] and the dtype is float. + void* mPtrMultiMemC{nullptr}; + + // The scaling factors calculated when quantizing C, for MxFp{4,8} and NvFp4 formats, also + // used for the DeepSeek FP8 recipe. + // + // For DeepSeek FP8 recipe: + // If transposeMmaOutput is false, shape is [N / 128, M]. + // Otherwise, shape is [M / 128, N]. + // The rightmost dimension is contiguous in memory. + // + // For MxFp{4,8} and NvFp4 formats: + // If transposeMmaOutput is false, shape is [M, N / 16]. + // Otherwise, shape is [N, M / 16]. + // The layout is controlled by options.mSfLayoutC (either R128c4 or R8c4). + // The layout (R128c4 and R8c4) is the same as explained in mPtrSfB. + // + // Otherwise should be set to nullptr. + void* mPtrSfC{nullptr}; + }; + + struct AllReduceBuffers { + // The barriers in global memory. + // + // The kernel arrives at (with release ordering) the multicast mapping of the barrier to + // broadcast amongst peer devices. It then waits (with acquire ordering) for the unicast mapping + // of the barrier. + // + // Flags in global memory that sync on "entrance" of reduce-scatter phase in two-shot + // all-reduce. The shape is [numTilesM * numTilesN] and the dtype is uint32_t. The pointer to + // the unicast memory created with IpcNvlsHandle. Must be set to 0 before the kernel launch. + void* mPtrTileBars{nullptr}; + + // The shape is [numTilesM * numTilesN] and the dtype is uint32_t. + // The pointer to the multicast memory created with IpcNvlsHandle. + void* mPtrMultiMemTileBars{nullptr}; + + // Flags in global memory that sync on "exit" after the all-reduce finishes. + // The shape is [numTilesM * numTilesN] and the dtype is uint32_t. + // The pointer to the unicast memory created with IpcNvlsHandle. + // Must be set to 0 before the kernel launch. + void* mPtrCompletionBars{nullptr}; + + // The shape is [numTilesM * numTilesN] and the dtype is uint32_t. + // The pointer to the multicast memory created with IpcNvlsHandle. + void* mPtrMultiMemCompletionBars{nullptr}; + }; + + ProblemDimensions mProblemDimensions; + InputBuffers mInputBuffers; + OutputBuffers mOutputBuffers; + AllReduceBuffers mAllReduceBuffers; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// +// +// GemmInterface +// +//////////////////////////////////////////////////////////////////////////////////////////////////// + +class GemmInterface { + public: + using ModuleCache = std::unordered_map>; + + GemmInterface() {} + + // Launch the cubin from the provided config. It calls all necessary memsets for internal buffers. + // Provided config must be validated with isValidConfig before the call. + int32_t run(GemmConfig const& config, void* workspace, GemmData const& options, void* cudaStream, + int32_t multiProcessorCount, bool usePdl = true, + std::optional> moduleCache = std::nullopt) const; + + // Initializes the buffers before the world sync. Must be called before run. + int32_t runInitBeforeWorldSync(GemmConfig const& config, GemmData const& data, + void* cudaStream) const; + + // Returns the size of the workspace buffers in bytes + size_t getWorkspaceSizeInBytes(GemmConfig const& config, GemmData const& data) const; + + // Returns the list of all available cubin configurations + GemmConfig const* getGemmConfigs() const; + + // Returns the number of available cubin configurations + size_t getNumGemmConfigs() const; + + // Returns true if the configuration of the cubin can be executed for the given params. + bool isValidConfig(GemmConfig const& config, GemmData const& data) const; + + private: + // Aligns the pointer to the alignment + template + inline Dtype* alignPtr(Dtype* ptr, int64_t alignment) const; + + // Creates GemmOptions from kernel and data. + GemmOptions getOptionsFromConfigAndData(GemmConfig const& config, GemmData const& data) const; + + // Returns the size of the workspace buffers in bytes + std::vector getWorkspaceSizesInBytes(GemmConfig const& config, + GemmData const& data) const; + + // Returns the size padded to the alignment + size_t getSizePaddedToAlignment(size_t size, size_t alignment) const; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +inline Dtype* GemmInterface::alignPtr(Dtype* ptr, int64_t alignment) const { + assert((alignment & (alignment - 1)) == 0 && "Alignment must be a power of 2"); + return reinterpret_cast((reinterpret_cast(ptr) + alignment - 1) & + ~(alignment - 1)); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +GemmConfig const* GemmInterface::getGemmConfigs() const { +#ifdef TLLM_GEN_EXPORT_INTERFACE + return tensorrt_llm::kernels::tllmGenGemmList; +#else + return nullptr; +#endif +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +size_t GemmInterface::getNumGemmConfigs() const { +#ifdef TLLM_GEN_EXPORT_INTERFACE + return tensorrt_llm::kernels::tllmGenGemmListLen; +#else + return 0; +#endif +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +GemmOptions GemmInterface::getOptionsFromConfigAndData(GemmConfig const& config, + GemmData const& data) const { + // Create options from config and data. + GemmOptions options; + options = config.mOptions; + options.mM = data.mProblemDimensions.mM; + options.mN = data.mProblemDimensions.mN; + options.mK = data.mProblemDimensions.mK; + return options; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +size_t GemmInterface::getSizePaddedToAlignment(size_t size, size_t alignment) const { + assert((alignment & (alignment - 1)) == 0); + return (size + alignment - 1) & ~(alignment - 1); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +size_t GemmInterface::getWorkspaceSizeInBytes(GemmConfig const& config, + GemmData const& data) const { + auto workspaceSizes = getWorkspaceSizesInBytes(config, data); + auto size = std::accumulate(workspaceSizes.begin(), workspaceSizes.end(), 0); + // Additional 1023 bytes to align the pointer to 1024 + return size > 0 ? size + 1023 : 0; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +std::vector GemmInterface::getWorkspaceSizesInBytes(GemmConfig const& config, + GemmData const& data) const { + // Get options from config. + auto& options = config.mOptions; + + // The number of tiles in the M dimension. + int32_t numTilesM = gemm::divUp(data.mProblemDimensions.mM, options.mTileM); + // The number of tiles in the N dimension. + int32_t numTilesN = gemm::divUp(data.mProblemDimensions.mN, options.mTileN); + + std::vector workspaceSizes; + + int64_t numBytesSplitK{0}, numBytesSplitKBars{0}; + if (doesSplitKUseGmem(options.mSplitK)) { + // The number of elements for intermediate split-k buffer that contains K slices padded to + // TileM/TileN sizes to avoid OOB accesses during the reduction. + // FIXME: Split-K has excessive memory traffic when combined with slice-K. + // Currently, data for all slice-K slices is sent, even though the slice-K reduction + // has already been performed. + // This should be optimized to send data for only one reduced slice. + auto numEltsSplitK = options.mNumSlicesForSplitK * numTilesM * numTilesN * options.mTileM * + options.mTileN * options.mNumSlicesForSliceK; + + // The number of bytes for intermediate split-k buffer that contains K slices. + numBytesSplitK = numEltsSplitK * tg::dtypeGetNumBits(tg::Dtype::Fp32) / /* bits */ 8; + // The number of bytes for the split-k completion barriers. + numBytesSplitKBars = numTilesM * numTilesN * sizeof(uint32_t); + // Two epilogue warps do GMEM split-k in DS GEMM. + if (options.mUseDeepSeekFp8) { + numBytesSplitKBars *= 2; + } + + // TODO: do we need to pad to 1024? + workspaceSizes.push_back(getSizePaddedToAlignment(numBytesSplitK, 1024)); + workspaceSizes.push_back(getSizePaddedToAlignment(numBytesSplitKBars, 1024)); + } + + return workspaceSizes; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +bool GemmInterface::isValidConfig(GemmConfig const& config, GemmData const& data) const { + // Get options from config and data. + auto options = getOptionsFromConfigAndData(config, data); + + // Is Blackwell? + bool isBlackwell = config.mSm == SmVersion::Sm100a; + + // Check options without modifications. + return checkAndUpdateGemmOptions(options, isBlackwell, data.mProblemDimensions.mWorldSize, + /* updateOptions */ false); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +int32_t GemmInterface::run(GemmConfig const& config, void* workspace, GemmData const& data, + void* cudaStream, int32_t multiProcessorCount, bool usePdl, + std::optional> moduleCache) const { + // Might be used. + (void)usePdl; + (void)moduleCache; + // Get options from config and data. + auto options = getOptionsFromConfigAndData(config, data); + + auto workspaceSizes = getWorkspaceSizesInBytes(config, data); + void* dSplitKSlices{nullptr}; + void* dPtrSplitKCompletionBars{nullptr}; + + // Set the completion barriers to 0 if needed. + if (doesSplitKUseGmem(options.mSplitK)) { + dSplitKSlices = alignPtr(reinterpret_cast(workspace), 1024); + dPtrSplitKCompletionBars = reinterpret_cast( + alignPtr(reinterpret_cast(dSplitKSlices) + workspaceSizes[0], 1024)); + auto err = cudaMemsetAsync((void*)dPtrSplitKCompletionBars, 0x00, workspaceSizes[1], + reinterpret_cast(cudaStream)); + if (err != cudaSuccess) { + return 1; + } + } + + // The number of tiles in the M dimension. + int numTilesM = gemm::divUp(options.mM, options.mTileM); + // The number of tiles in the N dimension. + int numTilesN = gemm::divUp(options.mN, options.mTileN); + + // Create kernel params. + auto kernelParams = gemm::KernelParams::setKernelParams( + options, data.mInputBuffers.mPtrA, data.mInputBuffers.mPtrSfA, + data.mInputBuffers.mPtrPerTokenSfA, data.mInputBuffers.mPtrB, data.mInputBuffers.mPtrSfB, + data.mInputBuffers.mPtrPerTokenSfB, data.mInputBuffers.mPtrBias, data.mOutputBuffers.mPtrC, + data.mOutputBuffers.mPtrSfC, data.mOutputBuffers.mPtrMultiMemC, + (float*)data.mInputBuffers.mPtrScaleC, dSplitKSlices, data.mAllReduceBuffers.mPtrTileBars, + data.mAllReduceBuffers.mPtrMultiMemTileBars, data.mAllReduceBuffers.mPtrCompletionBars, + data.mAllReduceBuffers.mPtrMultiMemCompletionBars, dPtrSplitKCompletionBars, + /* dPtrNumNonExitingCtas */ nullptr, data.mProblemDimensions.mRank, + data.mProblemDimensions.mWorldSize); + + // The size of the grid. + std::vector grid{numTilesM, numTilesN, options.mNumSlicesForSplitK}; + + // When split-k is enabled and to guarantee the forward progress, we must ensure that the number + // of tiles is less than number of SMs. This way, at least one CTA in the grid can make forward. + if (doesSplitKUseGmem(options.mSplitK)) { + if (grid[0] * grid[1] >= multiProcessorCount) { + // The number of MN tiles in Split-K (grid[0] * grid[1]) must be less than the number of SMs. + return 2; + } + } + +#ifdef TLLM_GEN_EXPORT_INTERFACE + CUmodule cuModule; + CUfunction cuFunction; + + auto fiModuleLoadData = [&](CUmodule* module) { + const std::string sha256 = config.mHash ? config.mHash : ""; + std::string fname_cubin = config.mFunctionName; + if (!fname_cubin.empty()) { + fname_cubin[0] = static_cast(std::toupper(static_cast(fname_cubin[0]))); + } + fname_cubin = tllm_gen_gemm_cubin_path + fname_cubin; + std::string cubin = flashinfer::trtllm_cubin_loader::getCubin(fname_cubin, sha256); + cuModuleLoadData(&cuModule, cubin.c_str()); + }; + if (moduleCache.has_value()) { + ModuleCache& moduleCacheRef = moduleCache.value().get(); + + // Modules are associated with a specific context, so the context is included in the key + CUcontext ctx; + unsigned long long ctxId; + cuCtxGetCurrent(&ctx); + cuCtxGetId(ctx, &ctxId); + + // Reinterpret the ctxId as a string to avoid needing a custom hash or converting it to a + // string in decimal representation. + std::string const ctxName = + std::string(reinterpret_cast(&ctxId), sizeof(unsigned long long) / sizeof(char)); + std::string const funcName = std::string(config.mFunctionName); + auto const moduleKey = ctxName + funcName; + auto module = moduleCacheRef.find(moduleKey); + + // Use cache if module is found, otherwise load and insert into cache + if (module != moduleCacheRef.end()) { + cuFunction = std::get<1>(module->second); + } else { + fiModuleLoadData(&cuModule); + cuModuleGetFunction(&cuFunction, cuModule, config.mFunctionName); + moduleCacheRef.insert(std::make_pair(moduleKey, std::make_tuple(cuModule, cuFunction))); + } + } else { + fiModuleLoadData(&cuModule); + cuModuleGetFunction(&cuFunction, cuModule, config.mFunctionName); + } + + // Prepare the grid/block. + dim3 block3{static_cast(config.mNumThreadsPerCTA), static_cast(1), + static_cast(1)}; + dim3 grid3{(grid.size() > 0 ? static_cast(grid[0]) : 1u), + (grid.size() > 1 ? static_cast(grid[1]) : 1u), + (grid.size() > 2 ? static_cast(grid[2]) : 1u)}; + // Prepare the cluster size. + dim3 cluster3{static_cast(options.mClusterDimX), + static_cast(options.mClusterDimY), + static_cast(options.mClusterDimZ)}; + + // Run the kernel. + auto result = trtllm::gen::launchKernel( + (void*)&kernelParams, cudaStream, config.mSharedMemSize, cuFunction, block3, grid3, cluster3, + usePdl && (config.mOptions.mGridWaitForPrimaryEarlyExit | + config.mOptions.mGridWaitForPrimaryA | config.mOptions.mGridWaitForPrimaryB)); + // If a module cache has not been given, unload the module to avoid leaking + if (!moduleCache.has_value()) { + cuModuleUnload(cuModule); + } + if (result != CUDA_SUCCESS) { + return -1; + } +#else + config.mCudaRunner->run((void*)&kernelParams, (void*)cudaStream, grid); +#endif + + return 0; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace gemm + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace gemm diff --git a/include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/GemmOptions.h b/include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/GemmOptions.h new file mode 100644 index 0000000000000000000000000000000000000000..5a9ed1b40cc8d57d136daf07271bd08694b57106 --- /dev/null +++ b/include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/GemmOptions.h @@ -0,0 +1,1202 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & + * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include +#include + +#include "Enums.h" +#include "KernelParams.h" +#include "KernelTraits.h" +#include "trtllm/gen/DtypeDecl.h" +#include "trtllm/gen/MmaDecl.h" +#include "trtllm/gen/SfLayoutDecl.h" +#ifndef TLLM_GEN_EXPORT_INTERFACE +#include "trtllm/gen/CudaRunner.h" +#include "trtllm/gen/GenCtx.h" +#else +#include + +template +void printArgs(T first, Args... args) { +#ifdef TLLM_GEN_DEBUG + std::cout << first; + if constexpr (sizeof...(args) > 0) { + std::cout << " "; + printArgs(args...); + } +#endif +} + +#define TLLM_CHECK_ERROR(cond, ...) \ + if (!(cond)) { \ + printArgs(__VA_ARGS__); \ + return false; \ + } + +#define TLLM_LOG_ERROR(...) TLLM_CHECK_ERROR(false, __VA_ARGS__) + +#define TLLM_CHECK_ERROR_FMT(cond, ...) TLLM_CHECK_ERROR(cond, __VA_ARGS__) + +#define TLLM_CHECK_WARNING(cond, ...) \ + if (!(cond)) { \ + printArgs(__VA_ARGS__); \ + return false; \ + } + +#define TLLM_LOG_WARNING(...) TLLM_CHECK_WARNING(false, __VA_ARGS__) + +#define TLLM_LOG_INFO(...) TLLM_CHECK_WARNING(false, __VA_ARGS__) + +#endif + +namespace gemm { + +namespace gemm { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace tg = trtllm::gen; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// NOTE: when adding new parameters, please update the dumpOptions function and +// gemm_export_config.json for cubin export. +struct GemmOptions { +#ifndef TLLM_GEN_EXPORT_INTERFACE + // allow safely down-casting + virtual ~GemmOptions() = default; +#endif + + GemmOptions() = default; + GemmOptions(AllReduceAlgo allReduceAlgo, BiasType biasType, int blockK, int clusterDimX, + int clusterDimY, int clusterDimZ, tg::Dtype dtypeAcc, tg::Dtype dtypeA, + tg::Dtype dtypeB, tg::Dtype dtypeC, tg::Dtype dtypeMmaA, tg::Dtype dtypeMmaB, + bool enablesEarlyExit, bool enablesDelayedEarlyExit, bool enablesGlobalPtxKnobs, + int epilogueLdtmDps, int epilogueLdtmBits, int epilogueTileM, int epilogueTileN, + bool gridTriggerSecondaryA, bool gridTriggerSecondaryB, + bool gridWaitForPrimaryEarlyExit, bool gridWaitForPrimaryA, bool gridWaitForPrimaryB, + bool hoistLoadTaskInit, bool hoistMmaTaskTryWaits, int k, KernelTraits kernelTraits, + MatrixLayout layoutA, MatrixLayout layoutB, int m, int mmaK, tg::MmaKind mmaKind, + int mmaM, int mmaN, bool mockAllReduce, int n, int numSlicesForSplitK, + int numSlicesForSliceK, int numStages, int numStagesMma, + int numStagesMmaWithinWorkTile, int numStagesMmaAcrossWorkTile, int numStagesWorkId, + bool outputDebugTensors, bool patchF2fp, bool useShuffledMatrixA, bool sliceK, + SplitK splitK, bool transposeMmaOutput, int tileM, int tileN, int tileK, + bool useUnrollLoop2xForMma, bool useCustomMmaSchedule, + bool useHoistTryWaitForCustomMmaSchedule, bool useDeepSeekFp8, bool usePerTokenSfA, + bool usePerTokenSfB, bool useTmaStore, bool useTwoTmaLoadWarps, bool useTwoMmaWarps, + tg::SfLayout sfLayoutA, tg::SfLayout sfLayoutB, tg::SfLayout sfLayoutC, + int sfReshapeFactor, TileScheduler tileScheduler) + : mAllReduceAlgo{allReduceAlgo}, + mBiasType{biasType}, + mBlockK(blockK), + mClusterDimX{clusterDimX}, + mClusterDimY{clusterDimY}, + mClusterDimZ{clusterDimZ}, + mDtypeAcc{dtypeAcc}, + mDtypeA{dtypeA}, + mDtypeB{dtypeB}, + mDtypeC{dtypeC}, + mDtypeMmaA{dtypeMmaA}, + mDtypeMmaB{dtypeMmaB}, + mEnablesEarlyExit{enablesEarlyExit}, + mEnablesDelayedEarlyExit{enablesDelayedEarlyExit}, + mEnablesGlobalPtxKnobs{enablesGlobalPtxKnobs}, + mEpilogueLdtmDps{epilogueLdtmDps}, + mEpilogueLdtmBits{epilogueLdtmBits}, + mEpilogueTileM{epilogueTileM}, + mEpilogueTileN{epilogueTileN}, + mGridTriggerSecondaryA{gridTriggerSecondaryA}, + mGridTriggerSecondaryB{gridTriggerSecondaryB}, + mGridWaitForPrimaryEarlyExit{gridWaitForPrimaryEarlyExit}, + mGridWaitForPrimaryA{gridWaitForPrimaryA}, + mGridWaitForPrimaryB{gridWaitForPrimaryB}, + mHoistLoadTaskInit{hoistLoadTaskInit}, + mHoistMmaTaskTryWaits{hoistMmaTaskTryWaits}, + mK{k}, + mKernelTraits{kernelTraits}, + mLayoutA{layoutA}, + mLayoutB{layoutB}, + mM{m}, + mMmaK{mmaK}, + mMmaKind{mmaKind}, + mMmaM{mmaM}, + mMmaN{mmaN}, + mMockAllReduce{mockAllReduce}, + mN{n}, + mNumSlicesForSplitK{numSlicesForSplitK}, + mNumSlicesForSliceK{numSlicesForSliceK}, + mNumStages{numStages}, + mNumStagesMma{numStagesMma}, + mNumStagesMmaWithinWorkTile{numStagesMmaWithinWorkTile}, + mNumStagesMmaAcrossWorkTile{numStagesMmaAcrossWorkTile}, + mNumStagesWorkId{numStagesWorkId}, + mOutputDebugTensors{outputDebugTensors}, + mPatchF2fp{patchF2fp}, + mUseShuffledMatrixA{useShuffledMatrixA}, + mSliceK{sliceK}, + mSplitK{splitK}, + mTransposeMmaOutput{transposeMmaOutput}, + mTileM{tileM}, + mTileN{tileN}, + mTileK{tileK}, + mUseUnrollLoop2xForMma{useUnrollLoop2xForMma}, + mUseCustomMmaSchedule{useCustomMmaSchedule}, + mUseHoistTryWaitForCustomMmaSchedule{useHoistTryWaitForCustomMmaSchedule}, + mUseDeepSeekFp8{useDeepSeekFp8}, + mUsePerTokenSfA{usePerTokenSfA}, + mUsePerTokenSfB{usePerTokenSfB}, + mUseTmaStore{useTmaStore}, + mUseTwoTmaLoadWarps{useTwoTmaLoadWarps}, + mUseTwoMmaWarps{useTwoMmaWarps}, + mSfLayoutA{sfLayoutA}, + mSfLayoutB{sfLayoutB}, + mSfLayoutC{sfLayoutC}, + mSfReshapeFactor{sfReshapeFactor}, + mTileScheduler{tileScheduler} {} + + // The all-reduce algorithm. + AllReduceAlgo mAllReduceAlgo{AllReduceAlgo::None}; + // The type of bias. + BiasType mBiasType{BiasType::None}; + // Block size in the K dimension + int mBlockK{-1}; + // Cluster size in X dim. + int mClusterDimX{1}; + // Cluster size in Y dim. + int mClusterDimY{1}; + // Cluster size in Z dim. + int mClusterDimZ{1}; + // Data type of the accumulators. + tg::Dtype mDtypeAcc{tg::Dtype::Fp32}; + // Data type of the A matrix. + tg::Dtype mDtypeA{tg::Dtype::Fp16}; + // Data type of the B matrix. + tg::Dtype mDtypeB{tg::Dtype::Void}; + // Data type of the outputs. + tg::Dtype mDtypeC{tg::Dtype::Void}; + // Data type of the A matrix for the MMA, if different from the input type. + tg::Dtype mDtypeMmaA{tg::Dtype::Void}; + // Data type of the B matrix for the MMA, if different from the input type. + tg::Dtype mDtypeMmaB{tg::Dtype::Void}; + // Whether to enable early exit. + bool mEnablesEarlyExit{false}; + // Whether to enable delayed early exit to overlap + // numNonExitingCtas loading with the other instructions. + bool mEnablesDelayedEarlyExit{false}; + // Whether to enable the global PTX knobs for guiding the compiler optimizations. + bool mEnablesGlobalPtxKnobs{true}; + // The epilogue supports multiple LDTM shapes, although not every shape is applicable in every + // case. In particular: + // - On Hopper: must be 16dp256bit. + // - Transposed output: must be 16dp256bit. + // - Non-transposed output: + // - NvFp4 with fused activation: must be 32dp32bit. + // - Else it can be either 16dp256bit or 32dp32bit. + // The number of DP lanes in the epilogue LDTM. + int mEpilogueLdtmDps{16}; + // The number of bits in the epilogue LDTM. + int mEpilogueLdtmBits{256}; + // Tile size for the epilogue in M dimension. + int mEpilogueTileM{128}; + // Tile size for the epilogue in N dimension. + int mEpilogueTileN{32}; + // Whether load task A triggers the next grid. + bool mGridTriggerSecondaryA{false}; + // Whether load task B triggers the next grid. + bool mGridTriggerSecondaryB{false}; + // Whether the loads that check for an early exit should wait on a grid dependency. + bool mGridWaitForPrimaryEarlyExit{true}; + // Whether the load of A should wait on a grid dependency. + bool mGridWaitForPrimaryA{true}; + // Whether the load of B should wait on a grid dependency. + bool mGridWaitForPrimaryB{true}; + // Whether to hoist the initialization of the loading tasks. + bool mHoistLoadTaskInit{true}; + // Whether to hoist the mbarrier try_waits (e.g., mma.prodAcq, smemAb.consWait) in the MMA task. + bool mHoistMmaTaskTryWaits{false}; + // The K dimension of GEMM. + int mK{16 * 16}; + // Traits of the kernel. + KernelTraits mKernelTraits{}; + // Layout of A matrix + MatrixLayout mLayoutA{MatrixLayout::MajorK}; + // Layout of B matrix + MatrixLayout mLayoutB{MatrixLayout::MajorK}; + // The M dimension of GEMM. + int mM{128 * 2}; + // Size of the MMA instruction in the K dimension. + int mMmaK{16}; + // The kind of MMA instruction to use. + tg::MmaKind mMmaKind{tg::MmaKind::Auto}; + // Size of the MMA instruction in the M dimension. + int mMmaM{64}; + // Size of the MMA instruction in the N dimension. + int mMmaN{16}; + // Whether to mock all-reduce code for single-GPU debugging. + bool mMockAllReduce{false}; + // The N dimension of GEMM. + int mN{64 * 4}; + // Number of partitions along the K dimension. When mNumSlicesForSplitK > 1, + // the problem is distributed across several SMs, where each CTA works on its local K slice. + // Partial results are accumulated afterwards using either GMEM or DSMEM (in CGA) + // to exchange the data between CTAs. + int mNumSlicesForSplitK{1}; + // Number of slices for slice-K along K dimension. + int mNumSlicesForSliceK{1}; + // The depth of the mainloop pipeline. + int mNumStages{2}; + // The depth of the mma pipeline. Equals numStagesMmaWithinWorkTile * numStagesMmaAcrossWorkTile. + int mNumStagesMma{1}; + // The depth of the mma pipeline within work tile. Only GmemC classes with "WithAccInReg" suffix + // are allowed to be greater than 1. + int mNumStagesMmaWithinWorkTile{-1}; + // The depth of the mma pipeline across work tiles in the persistent loop. + int mNumStagesMmaAcrossWorkTile{-1}; + // The depth of the work id pipeline and the work throttle pipeline. + int mNumStagesWorkId{3}; + // Whether to output debug tensors. + bool mOutputDebugTensors{false}; + // Patch float conversions. + bool mPatchF2fp{false}; + // Reorder rows/cols in the A matrix for the better memory accesses in the M-major epilogue. + bool mUseShuffledMatrixA{false}; + // Slice-K implementation to use TileM dimension for TileK. + bool mSliceK{false}; + // The location of the exchange for split-K (it's None when split-K is disabled). + SplitK mSplitK{SplitK::None}; + // Save output of MMA in M-major format. + bool mTransposeMmaOutput{false}; + // M tile dimension of GEMM. + int mTileM{128}; + // N tile dimension of GEMM. + int mTileN{32}; + // K tile dimension of GEMM. + int mTileK{16}; + // Whether to unroll the loop by 2x. + bool mUseUnrollLoop2xForMma{true}; + // Use custom MMA schedule optimized for low-latency. + bool mUseCustomMmaSchedule{false}; + // The purpose of hoisting trywaits is to opportunistically peek at the availability of the next + // k-block. It benefits when the next k-block is already available and thus sustaining the + // momentum, but it adds latency to the first k-block for smaller k-loop. + bool mUseHoistTryWaitForCustomMmaSchedule{false}; + // Use DeepSeek Fp8. + bool mUseDeepSeekFp8{false}; + // Apply per-token scales from A + bool mUsePerTokenSfA{false}; + // Apply per-token scales from B + bool mUsePerTokenSfB{false}; + // Use TMA to store the result. + bool mUseTmaStore{true}; + // Use two different warps for A and B matrix load. + bool mUseTwoTmaLoadWarps{false}; + // Use two different warps for MMA tasks. Applicable only to DeepSeek FP8. + bool mUseTwoMmaWarps{false}; + // Scale factors layout for A. + tg::SfLayout mSfLayoutA{tg::SfLayout::R128c4}; + // Scale factors layout for B. + tg::SfLayout mSfLayoutB{tg::SfLayout::R128c4}; + // Scale factors layout for C. + tg::SfLayout mSfLayoutC{tg::SfLayout::R128c4}; + // Number of "repeats", i.e. reshaping factor, to fold hidden dimension into SfBlock dimension. + // As result, the hidden dimension of the SF tensor must be a multiple of NumRepeats * + // numEltsPerSf * 4. This reduces the problem shape space that the kernel is able to run. + // But it reduces the number of L2 requests under the hood and potentially improves perf. + // Applies to layout 8x4 only. + int mSfReshapeFactor{1}; + // Tile scheduler type. + TileScheduler mTileScheduler{TileScheduler::Static}; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +enum class SmVersion { Sm90a, Sm100a }; + +//////////////////////////////////////////////////////////////////////////////////////////////////// +// +// GemmConfig +// +//////////////////////////////////////////////////////////////////////////////////////////////////// + +struct GemmConfig { + // When TRT-LLM Gen is exported to the other frameworks, the TLLM_GEN_EXPORT_INTERFACE must be + // defined. In this case, the cubins will be loaded from the provided data and function name. + // Otherwise, the kernel will be loaded from the CudaRunner. +#ifdef TLLM_GEN_EXPORT_INTERFACE + uint8_t const* mData{nullptr}; + uint32_t const mSize{0}; + uint32_t const mSharedMemSize{0}; + char const* mFunctionName{nullptr}; + uint32_t const mNumThreadsPerCTA{0}; + char const* mHash{nullptr}; +#else + trtllm::gen::CudaRunner* mCudaRunner{nullptr}; +#endif + + GemmOptions mOptions{}; + SmVersion mSm{SmVersion::Sm100a}; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// Serialization helpers. +template +inline std::string toString(T e) { + return std::to_string(e); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +inline std::string toString(trtllm::gen::Dtype e) { + return trtllm::gen::dtypeToString(e); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +inline std::string toString(trtllm::gen::MmaKind e) { + return trtllm::gen::mmaKindToString(e); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline std::string dumpOptions(GemmOptions const& options) { + std::stringstream ss; + ss << "mAllReduceAlgo=" << "gemm::AllReduceAlgo(" << static_cast(options.mAllReduceAlgo) + << ")" << "," << std::endl; + ss << "mBiasType=" << "gemm::BiasType(" << static_cast(options.mBiasType) << ")" << "," + << std::endl; + ss << "mBlockK=" << options.mBlockK << "," << std::endl; + ss << "mClusterDimX=" << options.mClusterDimX << "," << std::endl; + ss << "mClusterDimY=" << options.mClusterDimY << "," << std::endl; + ss << "mClusterDimZ=" << options.mClusterDimZ << "," << std::endl; + ss << "mDtypeAcc=" << "trtllm::gen::Dtype(" << static_cast(options.mDtypeAcc) << ")" + << "," << std::endl; + ss << "mDtypeA=" << "trtllm::gen::Dtype(" << static_cast(options.mDtypeA) << ")" << "," + << std::endl; + ss << "mDtypeB=" << "trtllm::gen::Dtype(" << static_cast(options.mDtypeB) << ")" << "," + << std::endl; + ss << "mDtypeC=" << "trtllm::gen::Dtype(" << static_cast(options.mDtypeC) << ")" << "," + << std::endl; + ss << "mDtypeMmaA=" << "trtllm::gen::Dtype(" << static_cast(options.mDtypeMmaA) << ")" + << "," << std::endl; + ss << "mDtypeMmaB=" << "trtllm::gen::Dtype(" << static_cast(options.mDtypeMmaB) << ")" + << "," << std::endl; + ss << "mEnablesEarlyExit=" << options.mEnablesEarlyExit << "," << std::endl; + ss << "mEnablesDelayedEarlyExit=" << options.mEnablesDelayedEarlyExit << "," << std::endl; + ss << "mEnablesGlobalPtxKnobs=" << options.mEnablesGlobalPtxKnobs << "," << std::endl; + ss << "mEpilogueLdtmDps=" << options.mEpilogueLdtmDps << "," << std::endl; + ss << "mEpilogueLdtmBits=" << options.mEpilogueLdtmBits << "," << std::endl; + ss << "mEpilogueTileM=" << options.mEpilogueTileM << "," << std::endl; + ss << "mEpilogueTileN=" << options.mEpilogueTileN << "," << std::endl; + ss << "mGridTriggerSecondaryA=" << options.mGridTriggerSecondaryA << "," << std::endl; + ss << "mGridTriggerSecondaryB=" << options.mGridTriggerSecondaryB << "," << std::endl; + ss << "mGridWaitForPrimaryEarlyExit=" << options.mGridWaitForPrimaryEarlyExit << "," << std::endl; + ss << "mGridWaitForPrimaryA=" << options.mGridWaitForPrimaryA << "," << std::endl; + ss << "mGridWaitForPrimaryB=" << options.mGridWaitForPrimaryB << "," << std::endl; + ss << "mHoistLoadTaskInit=" << options.mHoistLoadTaskInit << "," << std::endl; + ss << "mHoistMmaTaskTryWaits=" << options.mHoistMmaTaskTryWaits << "," << std::endl; + ss << "mK=" << options.mK << "," << std::endl; + ss << "mKernelTraits={}" << "," << std::endl; + ss << "mLayoutA=gemm::MatrixLayout(" << static_cast(options.mLayoutA) << ")" << "," + << std::endl; + ss << "mLayoutB=gemm::MatrixLayout(" << static_cast(options.mLayoutB) << ")" << "," + << std::endl; + ss << "mM=" << options.mM << "," << std::endl; + ss << "mMmaK=" << options.mMmaK << "," << std::endl; + ss << "mMmaKind=" << "trtllm::gen::MmaKind(" << static_cast(options.mMmaKind) << ")" + << "," << std::endl; + ss << "mMmaM=" << options.mMmaM << "," << std::endl; + ss << "mMmaN=" << options.mMmaN << "," << std::endl; + ss << "mMockAllReduce=" << options.mMockAllReduce << "," << std::endl; + ss << "mN=" << options.mN << "," << std::endl; + ss << "mNumSlicesForSplitK=" << options.mNumSlicesForSplitK << "," << std::endl; + ss << "mNumSlicesForSliceK=" << options.mNumSlicesForSliceK << "," << std::endl; + ss << "mNumStages=" << options.mNumStages << "," << std::endl; + ss << "mNumStagesMma=" << options.mNumStagesMma << "," << std::endl; + ss << "mNumStagesMmaWithinWorkTile=" << options.mNumStagesMmaWithinWorkTile << "," << std::endl; + ss << "mNumStagesMmaAcrossWorkTile=" << options.mNumStagesMmaAcrossWorkTile << "," << std::endl; + ss << "mNumStagesWorkId=" << options.mNumStagesWorkId << "," << std::endl; + ss << "mOutputDebugTensors=" << options.mOutputDebugTensors << "," << std::endl; + ss << "mPatchF2fp=" << options.mPatchF2fp << "," << std::endl; + ss << "mUseShuffledMatrixA=" << options.mUseShuffledMatrixA << "," << std::endl; + ss << "mSliceK=" << options.mSliceK << "," << std::endl; + ss << "mSplitK=" << "gemm::SplitK(" << static_cast(options.mSplitK) << ")" << "," + << std::endl; + ss << "mTransposeMmaOutput=" << options.mTransposeMmaOutput << "," << std::endl; + ss << "mTileM=" << options.mTileM << "," << std::endl; + ss << "mTileN=" << options.mTileN << "," << std::endl; + ss << "mTileK=" << options.mTileK << "," << std::endl; + ss << "mUseUnrollLoop2xForMma=" << options.mUseUnrollLoop2xForMma << "," << std::endl; + ss << "mUseCustomMmaSchedule=" << options.mUseCustomMmaSchedule << "," << std::endl; + ss << "mUseHoistTryWaitForCustomMmaSchedule=" << options.mUseHoistTryWaitForCustomMmaSchedule + << "," << std::endl; + ss << "mUseDeepSeekFp8=" << options.mUseDeepSeekFp8 << "," << std::endl; + ss << "mUsePerTokenSfA=" << options.mUsePerTokenSfA << "," << std::endl; + ss << "mUsePerTokenSfB=" << options.mUsePerTokenSfB << "," << std::endl; + ss << "mUseTmaStore=" << options.mUseTmaStore << "," << std::endl; + ss << "mUseTwoTmaLoadWarps=" << options.mUseTwoTmaLoadWarps << "," << std::endl; + ss << "mUseTwoMmaWarps=" << options.mUseTwoMmaWarps << "," << std::endl; + ss << "mSfLayoutA=" << "trtllm::gen::SfLayout(" << static_cast(options.mSfLayoutA) << ")" + << "," << std::endl; + ss << "mSfLayoutB=" << "trtllm::gen::SfLayout(" << static_cast(options.mSfLayoutB) << ")" + << "," << std::endl; + ss << "mSfLayoutC=" << "trtllm::gen::SfLayout(" << static_cast(options.mSfLayoutC) << ")" + << "," << std::endl; + ss << "mSfReshapeFactor=" << options.mSfReshapeFactor << "," << std::endl; + ss << "mTileScheduler=" << "gemm::TileScheduler(" << static_cast(options.mTileScheduler) + << ")" << std::endl; + return ss.str(); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +inline T divUp(T a, T b) { + return (a + b - 1) / b; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +inline T divUpMul(T a, T b) { + return gemm::divUp(a, b) * b; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline int32_t getShuffleBlockSize(int epilogueTileM) { + int shuffleBlockSize = 16; + if (epilogueTileM % 128 == 0) { + shuffleBlockSize = 32; + } + return shuffleBlockSize; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// Check if the options are valid or not. +inline bool checkAndUpdateGemmOptions(GemmOptions& options, bool isBlackwell, int /* tpGrpSize */, + bool updateOptions = true) { + if (options.mDtypeB == tg::Dtype::Void) { + if (updateOptions) { + options.mDtypeB = options.mDtypeA; + } else { + return false; + } + } + + // If not specified, used the input dtypes as MMA dtypes (no cast required). + if (options.mDtypeMmaA == tg::Dtype::Void) { + if (updateOptions) { + options.mDtypeMmaA = options.mDtypeA; + } else { + return false; + } + } + if (options.mDtypeMmaB == tg::Dtype::Void) { + if (updateOptions) { + options.mDtypeMmaB = options.mDtypeB; + } else { + return false; + } + } + + // Check that the A cast is supported. + // Currently, we only support {MxFp4, NvFp4} -> Bf16. + TLLM_CHECK_ERROR( + (options.mDtypeA == options.mDtypeMmaA) || + ((options.mDtypeA == tg::Dtype::MxE2m1 || options.mDtypeA == tg::Dtype::E2m1) && + options.mDtypeMmaA == tg::Dtype::Bfloat16), + "Unsupported cast for A: ", tg::dtypeToString(options.mDtypeA), " -> ", + tg::dtypeToString(options.mDtypeMmaA)); + + // Check that the B cast is supported. + // Currently, we only support Fp8 -> MxFp8. + // TODO: add same support for A (no transpose) + TLLM_CHECK_ERROR( + (options.mDtypeB == options.mDtypeMmaB) || + (options.mDtypeB == tg::Dtype::E4m3 && options.mDtypeMmaB == tg::Dtype::MxE4m3), + "Unsupported cast for B: ", tg::dtypeToString(options.mDtypeB), " -> ", + tg::dtypeToString(options.mDtypeMmaB)); + + if (options.mDtypeA != options.mDtypeMmaA) { + TLLM_CHECK_ERROR(options.mTileM == 128, + "TileM must be 128 when casting the input matrix A before the MMA."); + } + + if (options.mPatchF2fp) { + TLLM_CHECK_ERROR( + options.mDtypeA == tg::Dtype::MxE2m1 && options.mDtypeMmaA == tg::Dtype::Bfloat16, + "PatchF2fp is only supported for MxFp4 to Bf16 casts."); + } + + // FIXME: We do not support different dtypes for A and B when not on Blackwell. + if (!isBlackwell) { + TLLM_CHECK_ERROR(options.mDtypeMmaA == options.mDtypeMmaB, + "For non-Blackwell, A and B must have the same dtype."); + } + + // Check that the different dtypes for A and B are supported by the tensor core + // kind::f8f6f4 + if (options.mDtypeMmaA == tg::Dtype::E4m3 || options.mDtypeMmaA == tg::Dtype::E2m1) { + TLLM_CHECK_ERROR(options.mDtypeMmaB == tg::Dtype::E4m3 || options.mDtypeMmaB == tg::Dtype::E2m1, + "For dtypeMmaA = E4m3/E2m1 A, dtypeMmaB must also be E4m3/E2m1."); + } + + // kind::mxf8f6f4 + if (options.mDtypeMmaA == tg::Dtype::MxE4m3 || options.mDtypeMmaA == tg::Dtype::MxE2m1) { + TLLM_CHECK_ERROR( + options.mDtypeMmaB == tg::Dtype::MxE4m3 || options.mDtypeMmaB == tg::Dtype::MxE2m1, + "For dtypeMmaA = MxE4m3 or MxE2m1, dtypeMmaB must also be MxE4m3 or MxE2m1."); + } + if (options.mDtypeMmaB == tg::Dtype::MxE4m3 || options.mDtypeMmaB == tg::Dtype::MxE2m1) { + TLLM_CHECK_ERROR( + options.mDtypeMmaA == tg::Dtype::MxE4m3 || options.mDtypeMmaA == tg::Dtype::MxE2m1, + "For dtypeMmaB = MxE4m3 or MxE2m1, dtypeMmaA must also be MxE4m3 or MxE2m1."); + } + + // kind::f16 + if (options.mDtypeMmaA == tg::Dtype::Fp16 || options.mDtypeMmaA == tg::Dtype::Bfloat16) { + TLLM_CHECK_ERROR(options.mDtypeMmaB == options.mDtypeMmaA, + "For dtypeMmaA = Fp16/Bfloat16, dtypeMmaB must be the same as dtypeMmaA."); + } + + // When one of the inputs needs to be cast, we must use two load warps. + if ((options.mDtypeMmaA != options.mDtypeA || options.mDtypeMmaB != options.mDtypeB) && + !options.mUseTwoTmaLoadWarps) { + TLLM_LOG_WARNING("Two TMA load warps must be enabled if any of the inputs needs to be cast."); + } + + // When different dtypes are used for A and B, we must use different tiles to do the loading. + // It is not strictly required, but current implementation of SmemAb requires that. + if (options.mDtypeA != options.mDtypeB) { + TLLM_CHECK_ERROR(options.mUseTwoTmaLoadWarps, + "Two TMA load warps must be enabled for different input types of A and B."); + } + + // Get the mma kind for the input types. + if (options.mMmaKind == tg::MmaKind::Auto) { + if (updateOptions) { + options.mMmaKind = dtypeGetMmaKind(options.mDtypeMmaA, options.mDtypeMmaB); + } else { + return false; + } + } + + if ((options.mMmaKind == tg::MmaKind::Fp8Fp6Fp4 || + options.mMmaKind == tg::MmaKind::MxFp8Fp6Fp4) && + options.mMmaK != 32) { + TLLM_LOG_WARNING("Unsupported MmaK (", options.mMmaK, + ") for MmaKind=", gemm::toString(options.mMmaKind), ". Setting MmaK to 32"); + if (updateOptions) { + options.mMmaK = 32; + options.mTileK = std::max(options.mMmaK, options.mTileK); + } else { + return false; + } + } + + // Check LDTM shape. + if (isBlackwell) { + TLLM_CHECK_ERROR((options.mEpilogueLdtmDps == 16 && options.mEpilogueLdtmBits == 256) || + (options.mEpilogueLdtmDps == 32 && options.mEpilogueLdtmBits == 32), + "Unsupported LDTM shape: ", options.mEpilogueLdtmDps, "dp", + options.mEpilogueLdtmBits, "bit."); + if (options.mEpilogueTileM == 64) { + TLLM_CHECK_ERROR(options.mEpilogueLdtmDps == 16, + "Unsupported LDTM shape for epilogueTileM=64: ", options.mEpilogueLdtmDps, + "dp", options.mEpilogueLdtmBits, "bit."); + } + if (options.mTransposeMmaOutput) { + // We can't use 32dp32bit LDTM for transposed outputs because we need each thread to own + // multiple consecutive output elements. + TLLM_CHECK_ERROR((options.mEpilogueLdtmDps == 16 && options.mEpilogueLdtmBits == 256), + "Only 16dp256bit LDTM is supported for transposed outputs."); + } + } else { + TLLM_CHECK_ERROR( + options.mEpilogueLdtmDps == 16 && options.mEpilogueLdtmBits == 256, + "Hopper does not use TMEM. The register layout corresponds to 16dp256bit. Got ", + options.mEpilogueLdtmDps, "dp", options.mEpilogueLdtmBits, "bit."); + } + + // Constraints for NvFp4 and MxFp8. + if ((options.mMmaKind == tg::MmaKind::MxFp4NvFp4 || + options.mMmaKind == tg::MmaKind::MxFp8Fp6Fp4 || options.mDtypeC == tg::Dtype::MxE4m3) && + options.mMmaM != 128) { + // MMA M must be 128 when the input uses block scaling, or when the output is an Mx format. + int newTileM = 128 * divUp(options.mTileM, 128); + TLLM_LOG_WARNING("Unsupported MmaM (", options.mMmaM, + ") for MmaKind=", gemm::toString(options.mMmaKind), + ". Setting MmaM to 128 and TileM to ", newTileM); + if (updateOptions) { + options.mMmaM = 128; + options.mTileM = newTileM; + } else { + return false; + } + } + if (options.mMmaKind == tg::MmaKind::MxFp4NvFp4 || options.mMmaKind == tg::MmaKind::MxFp8Fp6Fp4) { + TLLM_CHECK_ERROR(isBlackwell, "Block scaling is only supported on Blackwell"); + + int const mmaK = (options.mMmaKind == tg::MmaKind::MxFp4NvFp4) ? 64 : 32; + if (options.mMmaK != mmaK) { + int newTileK = mmaK * divUp(options.mTileK, mmaK); + TLLM_LOG_WARNING("Unsupported MmaK (", options.mMmaK, + ") for MmaKind=", gemm::toString(options.mMmaKind), ". Setting MmaK to ", + mmaK, " and TileK to ", newTileK); + if (updateOptions) { + options.mMmaK = mmaK; + options.mTileK = newTileK; + } else { + return false; + } + } + + // The MMA N may only be smaller than 64 if it is equal to the tile N. + TLLM_CHECK_ERROR(options.mMmaN >= 64 || options.mMmaN == options.mTileN, "MmaN (", + options.mMmaN, ") must be >= 64 or equal to TileN (", options.mTileN, ")"); + } + if (tg::dtypeIsBlockFmt(options.mDtypeA)) { + int numEltsPerSfA = tg::dtypeNumEltsPerSf(options.mDtypeA); + TLLM_CHECK_ERROR(options.mTileK % (4 * numEltsPerSfA) == 0, "TileK (", options.mTileK, + ") must be a multiple of ", (4 * numEltsPerSfA), " for typeA ", + gemm::toString(options.mDtypeA)); + auto const numEltsPerSfAInK = options.mK / numEltsPerSfA; + TLLM_CHECK_ERROR(numEltsPerSfAInK % 4 == 0, "K dimension of scaling factors for A (", + numEltsPerSfAInK, ") must be a multiple of 4"); + } + if (tg::dtypeIsBlockFmt(options.mDtypeB)) { + TLLM_CHECK_ERROR(options.mSfLayoutB == tg::SfLayout::R128c4 || + options.mSfLayoutB == tg::SfLayout::R8c4 || + options.mSfLayoutB == tg::SfLayout::Linear, + "Only the 128x4 and 8x4 SF layouts are supported for B, got ", + tg::sfLayoutToString(options.mSfLayoutB)); + + // TileN must be a multiple of the number of rows per SF tile. + int const numSfTileRowsB = options.mSfLayoutB == tg::SfLayout::R128c4 ? 128 : 8; + TLLM_CHECK_ERROR(options.mTileN % numSfTileRowsB == 0, "TileN (", options.mTileN, + ") must be a multiple of ", numSfTileRowsB, " for B SF layout ", + tg::sfLayoutToString(options.mSfLayoutB)); + + int numEltsPerSfB = tg::dtypeNumEltsPerSf(options.mDtypeB); + TLLM_CHECK_ERROR(options.mTileK % (4 * numEltsPerSfB) == 0, "TileK (", options.mTileK, + ") must be a multiple of ", (4 * numEltsPerSfB), " for typeB ", + gemm::toString(options.mDtypeB)); + auto const numEltsPerSfBInK = options.mK / numEltsPerSfB; + TLLM_CHECK_ERROR(numEltsPerSfBInK % 4 == 0, "K dimension of scaling factors for B (", + numEltsPerSfBInK, ") must be a multiple of 4"); + } + + int32_t padMultiplierA = 1; + int32_t padMultiplierB = 1; + if (options.mMmaKind == tg::MmaKind::MxFp8Fp6Fp4) { + if (options.mDtypeA == tg::Dtype::MxE2m1) { + padMultiplierA = 2; + } + if (options.mDtypeB == tg::Dtype::MxE2m1) { + padMultiplierB = 2; + } + } + TLLM_CHECK_ERROR( + (padMultiplierA * tg::dtypeGetNumBits(options.mDtypeA) * options.mK / 8) % 16 == 0, + "K dimension of A must be aligned to 16 bytes."); + TLLM_CHECK_ERROR( + (padMultiplierB * tg::dtypeGetNumBits(options.mDtypeB) * options.mK / 8) % 16 == 0, + "K dimension of B must be aligned to 16 bytes."); + + if (options.mDtypeC == tg::Dtype::E2m1 || options.mDtypeC == tg::Dtype::MxE4m3) { + TLLM_CHECK_ERROR(isBlackwell, "Block scaling is only supported on Blackwell"); + + TLLM_CHECK_ERROR( + options.mSfLayoutC == tg::SfLayout::R128c4 || options.mSfLayoutC == tg::SfLayout::R8c4, + "Only the 128x4 and 8x4 SF layouts are supported for C."); + int const numSfTileRowsC = options.mSfLayoutC == tg::SfLayout::R128c4 ? 128 : 8; + int const tileTokenDim = options.mTransposeMmaOutput ? options.mTileN : options.mTileM; + TLLM_CHECK_ERROR_FMT(tileTokenDim % numSfTileRowsC == 0, + "Tile%s (%d) must be a multiple of %d for C SF layout %s", + options.mTransposeMmaOutput ? "N" : "M", tileTokenDim, numSfTileRowsC, + tg::sfLayoutToString(options.mSfLayoutC).c_str()); + + int const hiddenDim = options.mTransposeMmaOutput ? options.mM : options.mN; + int const hiddenGranularity = 4 * tg::dtypeNumEltsPerSf(options.mDtypeC); + TLLM_CHECK_ERROR(hiddenDim % hiddenGranularity == 0, "Hidden dim (", hiddenDim, + ") must be a multiple of ", hiddenGranularity, " for block-scaled outputs."); + TLLM_CHECK_ERROR(!options.mTransposeMmaOutput || options.mUseShuffledMatrixA, + "Transposing block-scaled outputs requires shuffled A."); + } + + // If dtypeC is unspecified (Dtype::Void), assign to the input dtype. + if (options.mDtypeC == tg::Dtype::Void) { + TLLM_LOG_INFO("Setting dtypeC to ", tg::dtypeToString(options.mDtypeA)); + if (updateOptions) { + options.mDtypeC = options.mDtypeA; + } else { + return false; + } + } + + // Set epilogue tile sizes to the output tile sizes, when epilogue tile sizes are incorrect. + if (options.mTileM % options.mEpilogueTileM != 0) { + TLLM_LOG_WARNING("TileM (", options.mTileM, ") must be divisible by EpilogueTileM (", + options.mEpilogueTileM, "). Setting EpilogueTileM to TileM"); + if (updateOptions) { + options.mEpilogueTileM = options.mTileM; + } else { + return false; + } + } + + if (options.mTileN % options.mEpilogueTileN != 0) { + TLLM_LOG_WARNING("TileN (", options.mTileN, ") must be divisible by EpilogueTileN (", + options.mEpilogueTileN, "). Setting EpilogueTileN to TileN"); + if (updateOptions) { + options.mEpilogueTileN = options.mTileN; + } else { + return false; + } + } + + // On Hopper, epilogue tile sizes are the same as output tiles. + if (!isBlackwell && + (options.mEpilogueTileM != options.mTileM || options.mEpilogueTileN != options.mTileN)) { + TLLM_LOG_WARNING( + "Overwriting epilogueTileM and epilogueTileN to match tileM and tileN respectively"); + if (updateOptions) { + options.mEpilogueTileM = options.mTileM; + options.mEpilogueTileN = options.mTileN; + } else { + return false; + } + } + + // Unsupported epilogue tile size. + if (options.mMmaM == 128 && options.mEpilogueTileM != options.mTileM) { + TLLM_LOG_WARNING( + "When MmaM = 128, EpilogueTileM must be equal to TileM. Setting EpilogueTileM to TileM"); + if (updateOptions) { + options.mEpilogueTileM = options.mTileM; + } else { + return false; + } + } + + TLLM_CHECK_ERROR(options.mM > 0 && options.mN > 0 && options.mK > 0, + "M, N and K must be larger than 0"); + TLLM_CHECK_ERROR(options.mNumSlicesForSplitK > 0, "Split K must be larger than 0."); + + if (options.mUseShuffledMatrixA) { + auto const shuffleBlockSize = getShuffleBlockSize(options.mEpilogueTileM); + TLLM_CHECK_ERROR(options.mM % shuffleBlockSize == 0, + "M must be a multiple of shuffle block size (", shuffleBlockSize, + ") when useShuffledMatrixA"); + } + + if (!options.mSliceK) { + TLLM_CHECK_ERROR(options.mMmaM <= options.mEpilogueTileM, + "EpilogueTileM must be larger or equal than mmaM."); + } + TLLM_CHECK_ERROR( + options.mTileM % options.mEpilogueTileM == 0 && options.mTileN % options.mEpilogueTileN == 0, + "TileM and TileN must be divisible by EpilogueTileM and EpilogueTileN respectively."); + TLLM_CHECK_ERROR(options.mClusterDimX == 1 && options.mClusterDimY == 1, + "GEMM does not support cluster in X and Y dimensions."); + TLLM_CHECK_ERROR(options.mClusterDimZ == 1 || options.mNumSlicesForSplitK > 1, + "Cluster DimZ is only allowed for split-k."); + TLLM_CHECK_ERROR(options.mTileM <= 128, "GEMM does not support TileM > 128."); + + // When the A-matrix is shuffled, the output must be transposed. + if (options.mUseShuffledMatrixA) { + // TODO add matrix shuffle for N-major epilogue. + TLLM_CHECK_ERROR( + options.mTransposeMmaOutput, + "Shuffled matrix A is only supported with M-major epilogue. Set -transposeMmaOutput"); + } + + // Check all-reduce options. + if (options.mAllReduceAlgo == AllReduceAlgo::OneShot) { + // One shot is implemented with PTX cp.reduce.async.bulk.tensor which supports only the + // following types for reduce add: u32, s32, u64, f32, f16, bf16. + // + // See: https://docs.nvidia.com/cuda/parallel-thread-execution/ + // #data-movement-and-conversion-instructions-cp-reduce-async-bulk-tensor + std::set dtypeSupported{tg::Dtype::UInt32, tg::Dtype::Int32, tg::Dtype::UInt64, + tg::Dtype::Fp32, tg::Dtype::Fp16, tg::Dtype::Bfloat16}; + TLLM_CHECK_ERROR(dtypeSupported.find(options.mDtypeC) != dtypeSupported.end(), + "Unsupported output dtype ", tg::dtypeToString(options.mDtypeC)); + } else if (options.mAllReduceAlgo == AllReduceAlgo::TwoShot) { + // TODO(anchengc): + // Input dtype == output dtype -> can perform all-reduce in-place. + // Input dtype != output dtype -> must perform all-reduce out of place. + TLLM_CHECK_ERROR_FMT( + options.mDtypeC == options.mDtypeAcc, + "Not implemented - mixed dtype (dtypeC (%s) != dtypeAcc (%s)) requires out of place update", + tg::dtypeToString(options.mDtypeC).c_str(), tg::dtypeToString(options.mDtypeAcc).c_str()); + } + if (options.mAllReduceAlgo != AllReduceAlgo::None) { + TLLM_CHECK_ERROR(options.mUseTmaStore, "Non-TMA store with all-reduce is not implemented"); + } + + if (updateOptions) { + if (options.mNumSlicesForSplitK == 1) { + // No split-k. + options.mSplitK = SplitK::None; + } else if (options.mNumSlicesForSplitK > 1 && options.mClusterDimZ == 1) { + // Split-k with exchange through gmem. + options.mSplitK = SplitK::Gmem; + } else { + // Split-k with exchange through Dsmem. + options.mSplitK = SplitK::Dsmem; + } + } + // For GMEM-based split-K, we write 4 elements at once. + if (options.mSplitK == SplitK::Gmem) { + TLLM_CHECK_ERROR((options.mM * options.mN) % 4 == 0, + "M * N must be a multiple of 4 for Split-K"); + } + + if (options.mNumSlicesForSplitK > 1) { + if ((options.mEpilogueTileM != options.mTileM || options.mEpilogueTileN != options.mTileN) && + !options.mUseDeepSeekFp8) { + TLLM_LOG_WARNING( + "Overwriting epilogueTileM and epilogueTileN to match tileM and tileN respectively"); + if (updateOptions) { + options.mEpilogueTileM = options.mTileM; + options.mEpilogueTileN = options.mTileN; + } else { + return false; + } + } + } + if (options.mSplitK == SplitK::Dsmem) { + TLLM_CHECK_ERROR(options.mClusterDimZ == options.mNumSlicesForSplitK, + "CGA size must be equal to the number of slices in split-k"); + } + + // Maps numStagesMma to (stagesWithinWorkTile, stagesAcrossWorkTile) if not already set. + // If (-1, -1) -> (numStagesMma / min(2, numStagesMma), min(2, numStagesMma)) + // If ( m, -1) -> (m, numStagesMma / m) + // If (-1, n) -> (numStagesMma / n, n) + if (options.mNumStagesMmaWithinWorkTile == -1 && options.mNumStagesMmaAcrossWorkTile == -1) { + if (updateOptions) { + options.mNumStagesMmaAcrossWorkTile = std::min(2, options.mNumStagesMma); + options.mNumStagesMmaWithinWorkTile = + options.mNumStagesMma / options.mNumStagesMmaAcrossWorkTile; + } else { + return false; + } + } else if (options.mNumStagesMmaWithinWorkTile == -1) { + if (updateOptions) { + options.mNumStagesMmaWithinWorkTile = + options.mNumStagesMma / options.mNumStagesMmaAcrossWorkTile; + } else { + return false; + } + } else if (options.mNumStagesMmaAcrossWorkTile == -1) { + if (updateOptions) { + options.mNumStagesMmaAcrossWorkTile = + options.mNumStagesMma / options.mNumStagesMmaWithinWorkTile; + } else { + return false; + } + } + // Check mma stages. + TLLM_CHECK_ERROR_FMT(options.mNumStagesMmaWithinWorkTile * options.mNumStagesMmaAcrossWorkTile == + options.mNumStagesMma && + options.mNumStagesMmaAcrossWorkTile <= 2, + "Condition numStagesMmaWithinWorkTile (%d) * numStagesMmaAcrossWorkTile " + "(%d) == numStagesMma (%d) && numStagesMmaAcrossWorkTile (%d) <= 2 must be " + "satisfied. Check arguments.", + options.mNumStagesMmaWithinWorkTile, options.mNumStagesMmaAcrossWorkTile, + options.mNumStagesMma, options.mNumStagesMmaAcrossWorkTile); + // Mma stage must be 1 for pre-Hopper. + TLLM_CHECK_ERROR(isBlackwell || options.mNumStagesMma == 1, + "Mma stage must be 1 for pre-Hopper. Found ", options.mNumStagesMma); + // DeepSeek Fp8 + if (!options.mUseDeepSeekFp8) { + TLLM_CHECK_ERROR(options.mNumStagesMmaWithinWorkTile == 1, + "Non-DeepSeekFp8 requires numStagesMmaWithinWorkTile == 1"); + if (options.mNumStagesMma > 1) { + TLLM_CHECK_ERROR(options.mTileScheduler == TileScheduler::Persistent, + "Non-DeepSeekFp8 requires persistent scheduler when using numStagesMma >1"); + } + } + if (options.mUseDeepSeekFp8) { + TLLM_CHECK_ERROR(options.mDtypeA == tg::Dtype::E4m3 && options.mDtypeB == tg::Dtype::E4m3, + "A and B dtype must be E4m3 for DeepSeek Fp8. Found dtypeA=", + tg::dtypeToString(options.mDtypeA), + " dtypeB=", tg::dtypeToString(options.mDtypeB)); + + TLLM_CHECK_ERROR(isBlackwell, "DeepSeek Fp8 is not supported for Hopper"); + TLLM_CHECK_ERROR(options.mAllReduceAlgo == AllReduceAlgo::None, + "DeepSeek Fp8 does not support AllReduce"); + + // Check that TileK = 128 for correct scaling of every 128 channels. + TLLM_CHECK_ERROR(options.mTileK == 128, "Tile-K must be equal to 128 for DeepSeek Fp8"); + TLLM_CHECK_ERROR(options.mK % options.mTileK == 0, "K must be a multiple of TileK"); + // Tile sizes of the output hidden dimension. + auto hiddenDimPerOutputTile = options.mTransposeMmaOutput ? options.mTileM : options.mTileN; + auto hiddenDimPerEpilogueTile = + options.mTransposeMmaOutput ? options.mEpilogueTileM : options.mEpilogueTileN; + auto hiddenDimPerMma = options.mTransposeMmaOutput ? options.mMmaM : options.mMmaN; + auto hiddenDimName = options.mTransposeMmaOutput ? "M" : "N"; + TLLM_CHECK_WARNING(options.mNumStagesMmaWithinWorkTile > 1, + "DeepSeekFp8 recommends setting \"-numStagesMmaWithinWorkTile 2\"."); + // Update the number of stages of the MMA accumulator pipeline. TODO: enable by default for + // deepseek. + // options.mNumStagesMma = 2; + // Use two MMA warps to reduce mbar trywait latency. TODO: enable by default for deepseek. + // options.mUseTwoMmaWarps = true; + + // Make sure the GEMM-K dimension is a multiple of 128 when using DeepSeek FP8. + TLLM_CHECK_ERROR(options.mK % 128 == 0, + "GEMM-K must be a multiple of 128 when using DeepSeek Fp8. Found ", + options.mK); + + // Check that the output tile N can be processed with the epilogue tile granularity. + TLLM_CHECK_ERROR((hiddenDimPerOutputTile / 2) % hiddenDimPerEpilogueTile == 0, + "DeepSeek Fp8 requires Tile", hiddenDimName, " / 2 (", + hiddenDimPerOutputTile / 2, ") being a multiple of EpilogueTile", + hiddenDimName, " (", hiddenDimPerEpilogueTile, ")"); + // Check that the output tile N can be processed with the epilogue tile granularity. + TLLM_CHECK_ERROR((hiddenDimPerOutputTile / 2) % hiddenDimPerMma == 0, + "DeepSeek Fp8 requires Tile", hiddenDimName, " / 2 (", + hiddenDimPerOutputTile / 2, ") being a multiple of mma", hiddenDimName, " (", + hiddenDimPerMma, ")"); + } + + if (options.mSliceK) { + TLLM_CHECK_ERROR(isBlackwell, "Slice-K is not supported on Hopper"); + + TLLM_CHECK_ERROR(!options.mUseDeepSeekFp8, "DeepSeek Fp8 GEMM is not supported for slice-K"); + TLLM_CHECK_ERROR(options.mUseTwoTmaLoadWarps, "Slice-K requires two warp load for A and B"); + TLLM_CHECK_ERROR(options.mTransposeMmaOutput, "Slice-K requires transpose mma output"); + TLLM_CHECK_ERROR(options.mUseShuffledMatrixA, "Slice-K requires shuffled matrix A"); + TLLM_CHECK_ERROR(options.mTileK % 128 == 0, "Slice-K requires TileK be a multiple of 128"); + TLLM_CHECK_ERROR(options.mMmaM == 128, "Slice-K requires MmaM == 128"); + TLLM_CHECK_ERROR(options.mTileN == options.mEpilogueTileN, + "TileN must be equal to EpilogueTileN for slice-K"); + + TLLM_LOG_WARNING("Overwriting TileM and EpilogueTileM to 32 for slice-K"); + if (options.mTileM != 32 || options.mEpilogueTileM != 32) { + if (updateOptions) { + // FIXME: it is possible to remove this restriction. + options.mTileM = 32; + options.mEpilogueTileM = 32; + } else { + return false; + } + } + TLLM_CHECK_ERROR(options.mDtypeA == tg::Dtype::E4m3 && options.mDtypeB == tg::Dtype::E4m3, + "Slice-K requires e4m3 input dtype"); + + if (options.mNumSlicesForSliceK != 4) { + if (updateOptions) { + options.mNumSlicesForSliceK = 4; + } else { + return false; + } + } + TLLM_CHECK_ERROR((options.mTileK / options.mMmaK) % options.mNumSlicesForSliceK == 0, "TileK (", + options.mTileK, ") / MmaK (", options.mMmaK, + ") must be a multiple of mNumSlicesForSliceK (", options.mNumSlicesForSliceK, + ")"); + } + + if (options.mUseUnrollLoop2xForMma) { + // Number of iterations in K dimension after padding. + // Note the perCtaK in each CTA in the splitK group are padded to the same number of iterations. + // E.g., K = 512, TileK = 128, numSlicesForSplitK = 3. Then the padded K is + // + // ceil(512 / (128*3)) * (128*3) = 768 + // + int paddedK = divUpMul(options.mK, options.mTileK * options.mNumSlicesForSplitK); + // Check that the padded K (K rounded to next multiple of tileK) is a multiple of 2*TileK when + // UnrollLoop2x is enabled. This is to avoid deadlock when mma runs even-numbered loop while the + // other warps run odd-numbered loop. + // + bool notSupported = (paddedK / options.mNumSlicesForSplitK) % (options.mTileK * 2) != 0; + if (notSupported) { + TLLM_LOG_WARNING("Size K / splitK must be a multiple of TileK * 2. Found TileK=", + options.mTileK, " and K=", options.mK, " (paddedK=", paddedK, + ") and numSlicesForSplitK=", options.mNumSlicesForSplitK, + ". Disabling unrollLoop2xForMma."); + if (updateOptions) { + options.mUseUnrollLoop2xForMma = false; + } else { + return false; + } + } + } + + if (!isBlackwell && options.mTileScheduler == TileScheduler::Persistent) { + // TODO(anchengc): will be supported in upcoming MRs. + TLLM_LOG_WARNING("Persistent scheduling is not supported on Hopper. Using Static scheduling."); + if (updateOptions) { + options.mTileScheduler = TileScheduler::Static; + } else { + return false; + } + } + + if (options.mEnablesDelayedEarlyExit && options.mEnablesEarlyExit) { + TLLM_LOG_WARNING( + "Only one of early exit and delayed early exit should be enabled. Disabling " + "delayed early exit"); + if (updateOptions) { + options.mEnablesDelayedEarlyExit = false; + } else { + return false; + } + } + + // This check prevents the triggering of the secondary (PREEXIT) from executing before the wait + // for primary (ACQBULK). This could lead to the following confusing situation, which we want to + // avoid: + // + // Kernel 3 is written with the assumption that it can read the output of + // kernel 1 *without* ACQBULK and the output of kernel 2 *with* ACQBULK. + // However, when we allow PREEXIT and ACQBULK to be executed out of order, + // this is not guaranteed. + // + // Time: ----> + // + // Kernel 1: ----PREEXIT-----------FLUSH + // Kernel 2: -------PREEXIT----ACQBULK---FLUSH + // Kernel 3: Warp 0: ---- (!) Output of 1,2 is not yet visible + // ----------------------- + // Warp 1: ---- (!) We normally assume that 1 is visible is not yet + // visible- Warp 2: -------------------ACQBULK-- Kernel 1,2 output visible + // ---------- + TLLM_CHECK_ERROR( + (options.mGridWaitForPrimaryA || !options.mGridTriggerSecondaryA), + "A: If a task triggers a secondary kernel, it must also wait for primary kernel."); + TLLM_CHECK_ERROR( + (options.mGridWaitForPrimaryB || !options.mGridTriggerSecondaryB), + "B: If a task triggers a secondary kernel, it must also wait for primary kernel."); + + if (options.mUsePerTokenSfA || options.mUsePerTokenSfB) { + // Checks applicable to both MetaFP8 and RoutingScalesOnInput + TLLM_CHECK_ERROR(!options.mUseDeepSeekFp8, + "DeepSeek FP8 and per-token scaling are not compatible"); + TLLM_CHECK_ERROR(isBlackwell, "Per-token scaling is not supported for Hopper"); + if (options.mUsePerTokenSfA && options.mUsePerTokenSfB) { + // MetaFP8 case + TLLM_CHECK_ERROR(options.mDtypeA == tg::Dtype::E4m3 && options.mDtypeB == tg::Dtype::E4m3, + "A and B dtype must be E4m3 for Meta Fp8. Found dtypeA=", + tg::dtypeToString(options.mDtypeA), + " dtypeB=", tg::dtypeToString(options.mDtypeB)); + } else { + // RoutingScalesOnInput case + TLLM_CHECK_ERROR((options.mUsePerTokenSfA && !options.mTransposeMmaOutput) || + (options.mUsePerTokenSfB && options.mTransposeMmaOutput), + "In RoutingScalesOnInput mode, perToken scales must be used on activations"); + } + } + + // The generation should support non K-major layouts for both A and B; however, it is unclear if + // there is a use-case + TLLM_CHECK_ERROR( + (options.mLayoutA == MatrixLayout::MajorK) || (options.mLayoutB == MatrixLayout::MajorK), + "At least one matrix must be in k-major layout"); + + // Some features are currently only support when both matrices are in K-major format + if (options.mLayoutB != MatrixLayout::MajorK || options.mLayoutB != MatrixLayout::MajorK) { + TLLM_CHECK_ERROR(isBlackwell, "Non K-major layouts are only supported on Blackwell"); + TLLM_CHECK_ERROR(options.mSplitK == SplitK::None, "Non K-major layouts do not support split K"); + } + if (options.mLayoutA == MatrixLayout::MajorMn) { + TLLM_CHECK_ERROR(tg::dtypeGetNumBits(options.mDtypeA) >= 8, + "Subbyte types only support K major layout"); + } + if (options.mLayoutB == MatrixLayout::MajorMn) { + TLLM_CHECK_ERROR(tg::dtypeGetNumBits(options.mDtypeB) >= 8, + "Subbyte types only support K major layout"); + } + + if ((options.mLayoutA == MatrixLayout::BlockMajorK) || + (options.mLayoutB == MatrixLayout::BlockMajorK)) { + bool const isBlockA = options.mLayoutA == MatrixLayout::BlockMajorK; + + // Block K size must be 128B. + // TODO Leaving this as an option for now in case we want to expertiment with other block sizes + // As the user is not expected to set this, do not fail if updateOptions is false + int32_t const elemSizeInBits = + (isBlockA) ? tg::dtypeGetNumBits(options.mDtypeA) : tg::dtypeGetNumBits(options.mDtypeB); + int32_t const elemsIn128B = 128 * 8 /* Bits in byte */ / elemSizeInBits; + + if (options.mBlockK != elemsIn128B) { + if (updateOptions) { + options.mBlockK = elemsIn128B; + } else { + return false; + } + } + + if (options.mBlockK > options.mTileK) { + TLLM_CHECK_ERROR( + options.mBlockK % options.mTileK == 0, + "If block size is greater than tile size, block size must be a multiple of tile size"); + } else if (options.mBlockK < options.mTileK) { + TLLM_CHECK_ERROR( + options.mTileK % options.mBlockK == 0, + "If tile size is greater than block size, tile size must be a multiple of block size"); + } + } + + if (!isBiasTypeNone(options.mBiasType)) { + TLLM_CHECK_ERROR(!isBiasTypeMn(options.mBiasType), "BiasType::Mn is not supported"); + TLLM_CHECK_ERROR(!options.mUseDeepSeekFp8, "Bias is not supported for DeepSeek Fp8"); + TLLM_CHECK_ERROR(!(options.mUsePerTokenSfA && options.mUsePerTokenSfB), + "Bias is not supported for Meta Fp8"); + } + + if (updateOptions) { + // Init kernel traits. + options.mKernelTraits = KernelTraits( + options.mDtypeA, options.mDtypeB, options.mDtypeC, options.mDtypeAcc, options.mDtypeMmaA, + options.mDtypeMmaB, options.mMmaKind, options.mTileM, options.mTileN, options.mTileK, + options.mEpilogueTileM, options.mEpilogueTileN, options.mNumStages, options.mNumStagesMma, + options.mNumSlicesForSplitK, options.mNumSlicesForSliceK, options.mSplitK, + options.mUseTmaStore, options.mTransposeMmaOutput, options.mAllReduceAlgo, + options.mTileScheduler == TileScheduler::Persistent, options.mUseDeepSeekFp8, + options.mUsePerTokenSfA, options.mUsePerTokenSfB, options.mBiasType); + } + + return true; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace gemm + +#ifdef TLLM_GEN_EXPORT_INTERFACE + +#undef TLLM_CHECK_ERROR +#undef TLLM_CHECK_ERROR_FMT +#undef TLLM_CHECK_WARNING +#undef TLLM_LOG_WARNING +#undef TLLM_LOG_INFO +#undef TLLM_LOG_ERROR + +#endif // TLLM_GEN_EXPORT_INTERFACE + +} // namespace gemm diff --git a/include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/KernelParams.h b/include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/KernelParams.h new file mode 100644 index 0000000000000000000000000000000000000000..0f1360d2a348f347ebe9a3f706cfca3436bc85cb --- /dev/null +++ b/include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/KernelParams.h @@ -0,0 +1,616 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & + * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include "Enums.h" +#include "TmaDescriptor.h" +#include "trtllm/gen/CommonUtils.h" +#include "trtllm/gen/SfLayoutDecl.h" + +namespace gemm { + +namespace gemm { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace tg = trtllm::gen; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +struct KernelParams { +#ifdef TLLM_ENABLE_CUDA + ////////////////////////////////////////////////////////////////////////////////////////////////// + // + // Gemm parameters. + // + ////////////////////////////////////////////////////////////////////////////////////////////////// + + // TMA descriptor for A. + // Must be setup using gemm::buildNdTmaDescriptor with shapes and strides from + // makeTmaShapeStrideAb. + // + // If layoutA is MatrixLayout::MajorK + // Logical shape is [M, K]. + // Logical strides are [K, 1]. + // Tile box shape is [tileM, tileK]. + // Tile box strides are [tileK, 1]. + // Dtype is set from options.mDtypeA. + // + // If layoutA is MatrixLayout::MajorMn + // Logical shape is [K, M]. + // Logical strides are [M, 1]. + // Tile box shape is [tileK, tileM]. + // Tile box strides are [tileM, 1]. + // Dtype is set from options.mDtypeA. + // + // If layoutA is MatrixLayout::BlockMajorK + // Logical shape is [K / blockK, M, blockK]. + // Logical strides are [M * blockK, blockK, 1]. + // Tile box shape is [tileK / min(blockK, tileK), tileM, min(blockK, tileK)]. + // Tile box strides are [tileM * min(blockK, tileK), min(blockK, tileK), 1]. + // Dtype is set from options.mDtypeA, and blockK is 128B. + CUtensorMap tmaA; + + // TMA descriptor for B. + // Must be setup using gemm::buildNdTmaDescriptor with shapes and strides from + // makeTmaShapeStrideAb. + // + // If layoutB is MatrixLayout::MajorK + // Logical shape is [N, K]. + // Logical strides are [K, 1]. + // Tile box shape is [tileN, tileK]. + // Tile box strides are [tileK, 1]. + // Dtype is set from options.mDtypeB. + // + // If layoutB is MatrixLayout::MajorMn + // Logical shape is [K, N]. + // Logical strides are [N, 1]. + // Tile box shape is [tileK, tileN]. + // Tile box strides are [tileN, 1]. + // Dtype is set from options.mDtypeB. + // + // If layoutB is MatrixLayout::BlockMajorK + // Logical shape is [K / blockK, N, blockK]. + // Logical strides are [N * blockK, blockK, 1]. + // Tile box shape is [tileK / min(blockK, tileK), tileN, min(blockK, tileK)]. + // Tile box strides are [tileN * min(blockK, tileK), min(blockK, tileK), 1]. + // Dtype is set from options.mDtypeB, and blockK is 128B. + CUtensorMap tmaB; + + // TMA descriptor for C, (when useTmaStore is true) + // Must be setup using gemm::buildNdTmaDescriptor with shapes and strides from + // makeTmaShapeStrideC. + // + // If transposeMmaOutput is false, + // Logical shape is [M, N]. + // Logical strides are [N, 1]. + // Tile box shape is [epilogueTileM, epilogueTileN]. + // Tile box strides are [epilogueTileN, 1]. + // Dtype is set from options.mDtypeC. + // + // If transposeMmaOutput is true, + // Logical shape is [N, M]. + // Logical strides are [M, 1]. + // Tile box shape is [epilogueTileN, epilogueTileM]. + // Tile box strides are [epilogueTileM, 1]. + // Dtype is set from options.mDtypeC. + CUtensorMap tmaC; + + // TMA descriptor for the block scaling factors for A, for MxFp{4,8} and NvFp4 formats. + // Must be setup using gemm::buildSfTmaDescriptor with shapes and strides from + // makeTmaShapeStrideSfAb. + // The layout of scaling factors for A is always R128c4 + // + // Let P be the number of elements per SF. P=16 for NvFp4, P=32 for Mx formats. + // K must be a multiple of 4P. + // The "logical" shape is: [M, K / P]. + // The R128c4 layout is: [⌈M / 128⌉, K / P / 4, 512]. + // The shape we use for TMA is: [⌈M / 128⌉, K / P / 4, 2, 256]. + // + // Dtype is Dtype::E4m3 for NvFp4, Dtype::UE8m0 for Mx formats. + CUtensorMap tmaSfA; + + // TMA descriptor for the block scaling factors for B, for MxFp{4,8} and NvFp4 formats. + // Must be setup using gemm::buildSfTmaDescriptor with shapes and strides from + // makeTmaShapeStrideSfAb. + // The layout of scaling factors for B is controlled by options.mSfLayoutB. + // + // Let P be the number of elements per SF. P=16 for NvFp4, P=32 for Mx formats. + // The "logical" shape is: [N, K / P] + // + // If the layout is R128c4, + // K must be a multiple of 4P. + // The R128c4 layout is: [⌈N / 128⌉, K / P / 4, 512] + // The shape we use for TMA is: [⌈N / 128⌉, K / P / 4, 2, 256] + // + // If the layout is R8c4, + // K must be a multiple of 4P. + // The R8c4 layout is: [⌈N / 8⌉, K / P / 4, 32] + // The shape we use for TMA is: [⌈N / 8⌉, K / P / 4 / r, r * 32] + // where r = min(tileK / P / 4, 8) + // + // Dtype is Dtype::E4m3 for NvFp4, Dtype::UE8m0 for Mx formats. + CUtensorMap tmaSfB; + + // The output matrix C. The data type is controlled by options.mDtypeC. + // + // When transposeMmaOutput is true, the shape is [N, M]. + // Otherwise, the shape is [M, N]. + // Elements in a given row are stored contiguously in memory (row-major). + void* ptrC; + + // The block scaling factors to dequantize A. + // + // If DeepSeek FP8 recipe is used: + // If transposeMmaOutput is false, shape is [K / 128, M]. + // Otherwise, shape is [M / 128, K / 128]. + // The rightmost dimension is contiguous in memory. + // + // If DeepSeek FP8 recipe is not used, but for MxFp{4,8} and NvFp4 formats: + // The layout and data type is the same as explained in tmaSfA. + // + // Otherwise should be set to nullptr. + void const* ptrSfA; + + // The scaling factors to dequantize B. + // + // If DeepSeek FP8 recipe is used: + // If transposeMmaOutput is false, shape is [N / 128, K / 128]. + // Otherwise, shape is [K / 128, N]. + // The rightmost dimension is contiguous in memory. + // + // If DeepSeek FP8 recipe is not used, but for MxFp{4,8} and NvFp4 formats: + // The layout and data type is the same as explained in tmaSfB. + // + // Otherwise should be set to nullptr. + void const* ptrSfB; + + // The bias applied after the GEMM. + // The bias is applied before applying the global scaling factor. I.e. + // C' = (A * B + bias') * scaleC + // scaleC = dequantA * dequantB * quantC + // Thus, the bias' = bias / (dequantA * dequantB), where the bias is the original bias. + // + // if BiasType is N, the shape is [N]. + // The bias is broadcasted along the M dimension. + // + // if BiasType is M, the shape is [M]. + // The bias is broadcasted along the N dimension. + // + // The dtype is float32. + void const* ptrBias; + + // The per-token scaling factors from scale A. + // + // This is used for either: + // * Per-token scaling factor quantization schemes, such as MetaFP8. The dtype is Dtype::Float32 + // * When the routing scales are applied to the input activations (only when output is not + // transposed). The dtype is Dtype::Bfloat16 + // + // The shape is [M] + void const* ptrPerTokenSfA; + + // The per-token scaling factors from scale B. + // + // This is used for either: + // * Per-token scaling factor quantization schemes, such as MetaFP8. The dtype is Dtype::Float32 + // * When the routing scales are applied to the input activations (only when output is + // transposed). The dtype is Dtype::Bfloat16 + // + // The shape is [N] + void const* ptrPerTokenSfB; + + // The scaling factors calculated when quantizing C, for MxFp{4,8} and NvFp4 formats, also + // used for the DeepSeek FP8 recipe. + // + // For DeepSeek FP8 recipe: + // If transposeMmaOutput is false, shape is [N / 128, M]. + // Otherwise, shape is [M / 128, N]. + // The rightmost dimension is contiguous in memory. + // + // For MxFp{4,8} and NvFp4 formats: + // If transposeMmaOutput is false, shape is [M, N / 16]. + // Otherwise, shape is [N, M / 16]. + // The layout is controlled by options.mSfLayoutC (either R128c4 or R8c4). + void* ptrSfC; + + // The output tensor scaling factor for MxFp{4,8}, Fp8, NvFp4 and DeepSeek FP8 quantization. + // TensorRT-LLM API requires a scaling factor on the device. + // Shape is [1]. + float const* ptrScaleC; + + // The M dimension. + // It is the total number of tokens if A is the activation matrix. + // It is the total number of output channels if A is the weight matrix. + int32_t m; + // The N dimension. + // It is the total number of tokens if B is the activation matrix. + // It is the total number of output channels if B is the weight matrix. + int32_t n; + // The K dimension. It is the hidden dimension of the input matrices. + int32_t k; + + ////////////////////////////////////////////////////////////////////////////////////////////////// + // + // All-reduce parameters. + // + ////////////////////////////////////////////////////////////////////////////////////////////////// + + // The rank id of the current device in the multi-gpu space. + int rank; + // The number of peer devices in tensor-parallel group. + int tpGrpSize; + // Pointer for output with multicast mapping. It is used by the "reduce" op (LDGMC.ADD) of the + // two-shot reduce-scatter phase. + // The shape is [M, N] and the dtype is float. + void* multimemC; + + // The barriers in global memory. + // + // The kernel arrives at (with release ordering) the multicast mapping of the barrier to broadcast + // amongst peer devices. It then waits (with acquire ordering) for the unicast mapping of the + // barrier. + // + // Flags in global memory that sync on "entrance" of reduce-scatter phase in two-shot all-reduce. + // The shape is [numTilesM * numTilesN] and the dtype is uint32_t. + // The pointer to the unicast memory created with IpcNvlsHandle. + // Must be set to 0 before the kernel launch. + void* ptrTileBars; + // The shape is [numTilesM * numTilesN] and the dtype is uint32_t. + // The pointer to the multicast memory created with IpcNvlsHandle. + void* multimemTileBars; + + // Flags in global memory that sync on "exit" after the all-reduce finishes. + // The shape is [numTilesM * numTilesN] and the dtype is uint32_t. + // The pointer to the unicast memory created with IpcNvlsHandle. + // Must be set to 0 before the kernel launch. + void* ptrCompletionBars; + // The shape is [numTilesM * numTilesN] and the dtype is uint32_t. + // The pointer to the multicast memory created with IpcNvlsHandle + void* multimemCompletionBars; + + ////////////////////////////////////////////////////////////////////////////////////////////////// + // + // Miscellaneous parameters. + // + ////////////////////////////////////////////////////////////////////////////////////////////////// + + // The barriers in global memory for Split-k reduction with exchange in GMEM. + // Each CTAs arrives at the barrier and blockIdx.z == gridDim.Z - 1 waits for the barrier to flip + // to perform a reduction. + // The shape is [numTilesM * numTilesN] and the dtype is uint32_t. + // For DeepSeek FP8 recipe, the shape is [numTilesM * numTilesN * 2]. + // The memory must be set to 0 before the kernel launch. + void* ptrSplitKCompletionBars; + + // Pointer to the memory holding the partial sums for split-K in GMEM. + // The shape is [numSlicesForSplitK, numSlicesForSliceK, numTilesM * tileM, numTilesN * tileN]. + // The dtype is dtypeAcc, i.e. float. + void* ptrPartialSumsForSplitK; + + // In some cases, some CTAs need to exit early. E.g. when the grid is statically set, but the + // actual workload is decided at runtime. This device pointer maps to the number of non exiting + // CTAs in the X dim of the grid when transposeMmaOutput is false. And the Y dim, otherwise. + // The pointer points to a scalar and the dtype is int32_t. The pointed value must be >= 0. + int32_t* ptrNumNonExitingCtas; + + ////////////////////////////////////////////////////////////////////////////////////////////////// + // + // Miscellaneous parameters. + // + ////////////////////////////////////////////////////////////////////////////////////////////////// + + enum class MatrixType { MatrixA = 0, MatrixB }; + + // Create the TMA shape/stride for A/B. + template + static auto makeTmaShapeStrideAb(GemmOptions const& options, MatrixType matrixType) { + // The outer dimension. + auto numTokens = (matrixType == MatrixType::MatrixA) ? options.mM : options.mN; + // The outer dimension tile size. + auto tileMn = (matrixType == MatrixType::MatrixA) ? options.mTileM : options.mTileN; + // The inner dimension. + auto hiddenSize = options.mK; + // The cute tensor shape for A/B: (numTokens, hiddenSize). + // Note that TMA descriptor expects the first dimension's stride to be + // 1, so swap the first two dimension so that the hiddenSize dimension comes first. + auto shape = + std::vector{static_cast(hiddenSize), static_cast(numTokens)}; + + // Assemble the stride (strideTokens, 1). + // Swap the first two dimension as mentioned before. + auto stride = std::vector{1, static_cast(hiddenSize)}; + + // Assemble the box shape + std::vector tileShape = {options.mTileK, tileMn}; + + MatrixLayout layout = (matrixType == MatrixType::MatrixA) ? options.mLayoutA : options.mLayoutB; + if (layout == MatrixLayout::MajorMn) { + // Apply transpose if necessary + std::swap(shape[0], shape[1]); + stride[1] = numTokens; + std::swap(tileShape[0], tileShape[1]); + } else if (layout == MatrixLayout::BlockMajorK) { + // Set shapes based on blocking layout + shape = {static_cast(options.mBlockK), static_cast(numTokens), + static_cast(options.mK / options.mBlockK)}; + stride = {1, static_cast(options.mBlockK), + static_cast(numTokens * options.mBlockK)}; + + // If blockK > tileK, then the inner most box size will be based on the tile + int32_t const tileBlockK = std::min(options.mBlockK, options.mTileK); + tileShape = {tileBlockK, tileMn, options.mTileK / tileBlockK}; + } + + return std::make_tuple(shape, stride, tileShape); + } + + // Create the TMA shape/stride for C. + template + static auto makeTmaShapeStrideC(GemmOptions const& options) { + // The number of tokens. + auto numTokens = options.mTransposeMmaOutput ? options.mN : options.mM; + // The hidden dimension. + auto hiddenSize = options.mTransposeMmaOutput ? options.mM : options.mN; + // Note that TMA descriptor expects the first dimension's stride to be + // 1, so swap the first two dimension so that the hiddenSize dimension comes first. + auto shape = + std::vector{static_cast(hiddenSize), static_cast(numTokens)}; + + // Assemble the stride (strideTokens, 1). + // Swap the first two dimension as mentioned before. + auto stride = std::vector{1, static_cast(hiddenSize)}; + + return std::make_tuple(shape, stride); + } + + // Create the TMA shape/stride for A/B block scaling factors. + template + static auto makeTmaShapeStrideSfAb(GemmOptions const& options, MatrixType matrixType, + tg::SfLayout layout) { + // The outer dimension. + auto numTokens = matrixType == MatrixType::MatrixA ? options.mM : options.mN; + // The inner dimension. + auto hiddenSize = options.mK; + // The outer tile dimension. + auto numTokensPerTile = matrixType == MatrixType::MatrixA ? options.mTileM : options.mTileN; + // The inner tile dimension. + auto hiddenSizePerTile = options.mTileK; + // The dtype of the matrix. + tg::Dtype matrixDtype = matrixType == MatrixType::MatrixA ? options.mDtypeA : options.mDtypeB; + // Number of elements per scaling factor. + int32_t const numEltsPerSf = (matrixDtype == tg::Dtype::E2m1) ? 16 : 32; + + switch (layout) { + case tg::SfLayout::R128c4: { + // The scaling factor tensor packs 128x4 tiles into contiguous 512B blocks. + // The 512B block maps to a 32x16B (32x128b) block in TMEM. + // See https://nvbugspro.nvidia.com/bug/4165523 + // + // Additionally, we have to meet constraints of TMA that the box dimensions are less + // than 256 and boxDim[0] is a multiple of 16B. + // + // The "logical" tensor is: [outer, inner / numEltsPerSf] + // The aforementioned format is: [⌈outer / 128⌉, inner / (4 * numEltsPerSf), 512] + // The shape we use for TMA is: [⌈outer / 128⌉, inner / (4 * numEltsPerSf), 2, 256] + + auto shape = std::vector{ + 256, 2, static_cast(tg::ceilDiv(hiddenSize, numEltsPerSf * 4)), + static_cast(tg::ceilDiv(numTokens, 128))}; + + std::vector stride(shape.size()); + stride[0] = 1; + for (size_t i = 1; i < shape.size(); i++) { + stride[i] = shape[i - 1] * stride[i - 1]; + } + + auto tileShapes = std::vector{ + 256, 2, static_cast(tg::ceilDiv(hiddenSizePerTile, numEltsPerSf * 4)), + static_cast(tg::ceilDiv(numTokensPerTile, 128))}; + + return std::make_tuple(shape, stride, tileShapes); + } + + case tg::SfLayout::R8c4: { + // The scaling factor tensor packs 8x4 tiles into contiguous 32B blocks. + // + // As the inner dimension (k) is often a multiple of the tile size, we can reshape to use + // fewer read requests, if the tile dimensions allow. It does not reduce the number of + // instructions. + // + // I.e., let's define r = min(⌈hiddenSizePerTile / (numEltsPerSf * 4)⌉, 8) + // + // The "logical" tensor is: [outer, inner / numEltsPerSf] + // The 8x4 SF layout is: [⌈outer / 8⌉, inner / (4 * numEltsPerSf), 32] + // The TMA tensor shape is: [⌈outer / 8⌉, inner / (4 * numEltsPerSf * r), r * 32] + // + // The caveat of NumRepeats>1 is we must pad the hidden dimension of SF to multiples of + // NumRepeats * numEltsPerSf * 4. + + // Detect if the supplied factor is power of 2. E.g., 0b0100 and (0b0100 - 1) == 0b0000. + int const r = options.mSfReshapeFactor; + if (r > 0 && (r & (r - 1)) != 0) { + throw std::runtime_error("mSfReshapeFactor must be positive and a power of 2. Found " + + std::to_string(r)); + } + + // Sanitize number of repeats so it doesn't exceed the dimension. + int const repeats = std::min(tg::ceilDiv(hiddenSizePerTile, numEltsPerSf * 4), r); + + // Detect if the input hidden size K is a multiple of the repeats. + if (tg::ceilDiv(hiddenSize, numEltsPerSf * 4) % repeats != 0) { + throw std::runtime_error( + "SF hiddenSize K (" + std::to_string(tg::ceilDiv(hiddenSize, numEltsPerSf * 4)) + + ") must be a multiple of repeats (" + std::to_string(repeats) + ")"); + } + + auto shape = std::vector{ + static_cast(repeats * 32), + static_cast(tg::ceilDiv(hiddenSize, numEltsPerSf * 4 * repeats)), + static_cast(tg::ceilDiv(numTokens, 8))}; + + std::vector stride(shape.size()); + stride[0] = 1; + for (size_t i = 1; i < shape.size(); i++) { + stride[i] = shape[i - 1] * stride[i - 1]; + } + + auto tileShapes = std::vector{ + static_cast(repeats * 32), + static_cast(tg::ceilDiv(hiddenSizePerTile, numEltsPerSf * 4 * repeats)), + static_cast(tg::ceilDiv(numTokensPerTile, 8))}; + + return std::make_tuple(shape, stride, tileShapes); + } + + default: + throw std::runtime_error("Unsupported SF layout"); + } + return std::make_tuple(std::vector{}, std::vector{}, + std::vector{}); + } + + // Setup the kernel parameters. + template + static KernelParams setKernelParams(GemmOptions_ const& options, void const* ptrA, + void const* ptrSfA, void const* ptrPerTokenSfA, + void const* ptrB, void const* ptrSfB, + void const* ptrPerTokenSfB, void const* ptrBias, void* ptrC, + void* ptrSfC, void* multimemC, float* ptrScaleC, + void* ptrPartialSumsForSplitK, void* ptrTileBars, + void* multimemTileBars, void* ptrCompletionBars, + void* multimemCompletionBars, void* ptrSplitKCompletionBars, + int32_t* ptrNumNonExitingCtas, int rank, int tpGrpSize) { + // Is one-shot all-reduce? + bool const oneShotAr{options.mAllReduceAlgo == AllReduceAlgo::OneShot}; + // Is two-shot all-reduce? + bool const twoShotAr{options.mAllReduceAlgo == AllReduceAlgo::TwoShot}; + // Are there peer devices? + bool const multiDevice{tpGrpSize > 1}; + + // Create the return struct. + KernelParams params; + + // Shape/stride for gmem tensor A. + auto [shapeA, strideA, tileShapeA] = makeTmaShapeStrideAb(options, MatrixType::MatrixA); + // Build tma descriptor for A. + params.tmaA = gemm::buildNdTmaDescriptor(options.mDtypeA, options.mMmaKind, shapeA, strideA, + tileShapeA, const_cast(ptrA)); + + // Shape/stride for gmem tensor B. + auto [shapeB, strideB, tileShapeB] = makeTmaShapeStrideAb(options, MatrixType::MatrixB); + // Build tma descriptor for B. + params.tmaB = gemm::buildNdTmaDescriptor(options.mDtypeB, options.mMmaKind, shapeB, strideB, + tileShapeB, const_cast(ptrB), + /* swizzle */ !options.mSliceK); + + if (options.mDtypeA == tg::Dtype::E2m1 || options.mDtypeA == tg::Dtype::MxE2m1 || + options.mDtypeA == tg::Dtype::MxE4m3) { + tg::Dtype const dTypeSfA = + (options.mDtypeA == tg::Dtype::E2m1) ? tg::Dtype::E4m3 : tg::Dtype::UE8m0; + + // Build TMA descriptor for gmem A block scaling factors. + auto [shapeSfA, strideSfA, tileShapesSfA] = + makeTmaShapeStrideSfAb(options, MatrixType::MatrixA, tg::SfLayout::R128c4); + params.tmaSfA = gemm::buildSfTmaDescriptor(dTypeSfA, shapeSfA, strideSfA, tileShapesSfA, + const_cast(ptrSfA)); + } + + if (options.mDtypeB == tg::Dtype::E2m1 || options.mDtypeB == tg::Dtype::MxE2m1 || + options.mDtypeB == tg::Dtype::MxE4m3) { + tg::Dtype const dTypeSfB = + (options.mDtypeB == tg::Dtype::E2m1) ? tg::Dtype::E4m3 : tg::Dtype::UE8m0; + + // Build TMA descriptor for gmem B block scaling factors. + auto [shapeSfB, strideSfB, tileShapesSfB] = + makeTmaShapeStrideSfAb(options, MatrixType::MatrixB, options.mSfLayoutB); + params.tmaSfB = gemm::buildSfTmaDescriptor(dTypeSfB, shapeSfB, strideSfB, tileShapesSfB, + const_cast(ptrSfB)); + } + + if (options.mUseTmaStore) { + // Shape/stride for gmem tensor C. + auto [shapeC, strideC] = makeTmaShapeStrideC(options); + + // Swap M and N tiles for the M-major epilogue. + auto outputTileM = + options.mTransposeMmaOutput ? options.mEpilogueTileN : options.mEpilogueTileM; + auto outputTileN = + options.mTransposeMmaOutput ? options.mEpilogueTileM : options.mEpilogueTileN; + + // One-shot performs TMA reduction on multicast mapping of the output buffer directly. + // Two-shot performs TMA store on unicast mapping of the output buffer. The reduction happens + // in the next phase. + void* ptrTmaC{oneShotAr && multiDevice ? multimemC : ptrC}; + auto dtypeC{options.mDtypeC}; + // Regardless of output dtype, two-shot all-reduce store partial + // accumulation results to global memory in float32 precision. + if (twoShotAr && multiDevice) { + dtypeC = options.mDtypeAcc; + } + + // Build tma descriptor for C. + params.tmaC = gemm::buildNdTmaDescriptor(dtypeC, tg::MmaKind::Auto, shapeC, strideC, + std::vector{outputTileN, outputTileM}, + const_cast(ptrTmaC)); + } + + // Set the dequantization factors for A and B when DeepSeek FP8 recipe is used. + params.ptrSfA = ptrSfA; + params.ptrSfB = ptrSfB; + + // Set the per-token scale factors for MetaFP8 or scale inputs + params.ptrPerTokenSfA = ptrPerTokenSfA; + params.ptrPerTokenSfB = ptrPerTokenSfB; + + // Set the bias. + params.ptrBias = ptrBias; + + // Also set ptrC (it may be used by the NCCL reduction code in "layers/Llama"). + params.ptrC = ptrC; + params.ptrScaleC = ptrScaleC; + + // The block scaling factors of C for MxFp{4,8} and NvFp4 formats. + // (not to be confused with the tensor-level scaling factor stored in ptrScaleC) + params.ptrSfC = ptrSfC; + + params.m = options.mM; + params.n = options.mN; + params.k = options.mK; + + params.rank = rank; + params.tpGrpSize = tpGrpSize; + + params.multimemC = multimemC; + params.ptrPartialSumsForSplitK = ptrPartialSumsForSplitK; + params.ptrTileBars = ptrTileBars; + params.multimemTileBars = multimemTileBars; + params.ptrCompletionBars = ptrCompletionBars; + params.multimemCompletionBars = multimemCompletionBars; + + params.ptrSplitKCompletionBars = ptrSplitKCompletionBars; + params.ptrNumNonExitingCtas = ptrNumNonExitingCtas; + return params; + } +#endif +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace gemm + +} // namespace gemm diff --git a/include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/KernelTraits.h b/include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/KernelTraits.h new file mode 100644 index 0000000000000000000000000000000000000000..9c3211f9427446008bcad5765ffb494451a36f6c --- /dev/null +++ b/include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/KernelTraits.h @@ -0,0 +1,619 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & + * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include + +#include "Enums.h" +#include "trtllm/gen/CommonUtils.h" +#include "trtllm/gen/DtypeDecl.h" + +namespace gemm { + +namespace gemm { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace tg = trtllm::gen; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// Structure to manage memory allocation with configurable reuse +class MemAllocatorHelper { + public: + // The default constructor. + MemAllocatorHelper() {} + + // Constructor to initialize chunk sizes, alignments, and reuse flags + MemAllocatorHelper(std::vector> const& sizes, + std::vector const& reuse, std::vector const& names) + : mNumBytesAndAlignmentPerSmemChunk(sizes), mFirstChunkReuse(reuse), mSmemChunkNames(names) {} + + // Function to calculate the size of the array from 0 to jj chunks + int32_t getOffsetBeforeChunk(int jj) const { + int32_t totalSize = 0; + for (int32_t ii = 0; ii < jj; ++ii) { + auto const& elem = mNumBytesAndAlignmentPerSmemChunk[ii]; + auto paddedSize = getSizePaddedToAlignment(elem.first, elem.second); + // If SMEM chunk is reused but the size of the current chunk is + // larger than currently counted size + if (mFirstChunkReuse[ii] && paddedSize > totalSize) { + // Set new size to the size of the current chunk. + // E.g. possible in case of + // mNumBytesAndAlignmentPerSmemChunk = {{1, 1}, {1, 1}, {1024, 1}} + // mFirstChunkReuse = {false, false, true} + // The last chunk is larger than the first plus second, so total size is 1024. + totalSize = paddedSize; + } else if (!mFirstChunkReuse[ii]) { + totalSize += paddedSize; + } + } + return totalSize; + } + + // Returns the offset of the ith chunk + int32_t getChunkOffset(int32_t ii) const { + if (mFirstChunkReuse[ii]) { + // Reuse the offset of the 0th chunk. + return getChunkOffset(0); + } + + // Get offset of ii chunks. + auto offset = getOffsetBeforeChunk(ii); + // Ensure alignment for the current chunk + return getSizePaddedToAlignment(offset, mNumBytesAndAlignmentPerSmemChunk[ii].second); + } + + // Function to calculate the total size of the SMEM array + int32_t getTotalSize() const { + return getOffsetBeforeChunk(static_cast(mNumBytesAndAlignmentPerSmemChunk.size())); + } + + // Returns the first chunk reuse flag for the ith chunk. + int getFirstChunkReuseFlag(int32_t ii) const { return mFirstChunkReuse[ii]; } + + // Print the contents of this object. + void print() const { + for (size_t ii = 0; ii < mNumBytesAndAlignmentPerSmemChunk.size(); ++ii) { + printf("Chunk %zd %s: %d bytes, %d alignment, reuse %s, offset %d\n", ii, + mSmemChunkNames[ii].c_str(), mNumBytesAndAlignmentPerSmemChunk[ii].first, + mNumBytesAndAlignmentPerSmemChunk[ii].second, mFirstChunkReuse[ii] ? "true" : "false", + getChunkOffset(ii)); + } + } + + private: + // Helper function to calculate padded size + int32_t getSizePaddedToAlignment(int32_t size, int32_t alignment) const { + assert((alignment & (alignment - 1)) == 0); + return (size + alignment - 1) & ~(alignment - 1); + } + + private: + // Sizes and alignment requirements of each chunk + // NOTE: be careful and make sure that the memory dependency is clear and + // chunks in the beginning of the SMEM can be overwritten. + std::vector> mNumBytesAndAlignmentPerSmemChunk; + // Chunk reuse configuration. True at ith position means that ith chunk starts at smemOffset = 0. + std::vector mFirstChunkReuse; + // Buffer names for inspection purposes. + std::vector mSmemChunkNames; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +int getNumSmemBitsPerElt(tg::Dtype dtype, tg::MmaKind mmaKind) { + if (mmaKind == tg::MmaKind::Auto) { + std::cout << "mmaKind != tg::MmaKind::Auto" << std::endl; + assert(false); + return -1; + } + if (mmaKind == tg::MmaKind::MxFp8Fp6Fp4) { + return 8; + } else { + return tg::dtypeGetNumBits(dtype); + } +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +class KernelTraits { + public: + // The default constructor. + KernelTraits() {} + + // The constructor. + KernelTraits(tg::Dtype dtypeA, tg::Dtype dtypeB, tg::Dtype dtypeC, tg::Dtype dtypeAcc, + tg::Dtype dtypeMmaA, tg::Dtype dtypeMmaB, tg::MmaKind mmaKind, int32_t tileM, + int32_t tileN, int32_t tileK, int32_t epilogueTileM, int32_t epilogueTileN, + int32_t numStages, int32_t numStagesMma, int32_t numSlicesForSplitK, + int32_t numSlicesForSliceK, SplitK splitK, bool useTmaStore, bool transposeMmaOutput, + AllReduceAlgo allReduceAlgo, bool usePersistentScheduler, bool useDeepSeekFp8, + bool usePerTokenSfA, bool usePerTokenSfB, BiasType biasType) + : mMmaKind{mmaKind} { + // + // SMEM + // + { + // [smemA ] (1024B aligned) + // [smemB ] (1024B aligned) + // [smemBShuffle ] (1024B aligned) + // [gmemC0 ] (1024B aligned) (if needed) + // [gmemC1 ] (1024B aligned) (if needed) + // [rowMax ] (16B aligned) (if needed) + // [sliceK ] (16B aligned) (if needed) + // [per-token SF ] (16B aligned) (if needed) + // [bias ] (16B aligned) (if needed) + // + // SMEM for smemA and smemB might be repurposed and used for gmemC0 and gmemC1: + // + // [..smemA..][..smemB..][..smemBShuffle..] + // [..gmemC0..][..gmemC1..][..rowMax..][..sliceK..][..per-token SF..][..bias..] + // + + if (mMmaKind == tg::MmaKind::Auto) { + mMmaKind = dtypeGetMmaKind(dtypeMmaA, dtypeMmaB); + } + + std::vector> numBytesAndAlignmentPerSmemChunk; + std::vector firstChunkReuseSmem; + // Buffer names for inspection purposes. + std::vector smemChunkNames; + + // LoadA + { + // Number of bytes in load A shared memory. + auto const numSmemBytesLoadA = + numStages * tileM * tileK * getNumSmemBitsPerElt(dtypeA, mMmaKind) / 8 /* bits */; + // Number of bytes for load A alignment for TMA load. + auto const numBytesAlignmentLoadA = 1024; + // loadA is already at first chunk. No need to reuse it. + auto const reuseChunksSmemLoadA = false; + // Add info. + smemChunkNames.emplace_back("smemLoadA"); + numBytesAndAlignmentPerSmemChunk.emplace_back( + std::make_pair(numSmemBytesLoadA, numBytesAlignmentLoadA)); + firstChunkReuseSmem.emplace_back(reuseChunksSmemLoadA); + } + + // LoadB + { + // Number of bytes in load B shared memory. + auto const numSmemBytesLoadB = + numStages * tileN * tileK * getNumSmemBitsPerElt(dtypeB, mMmaKind) / 8 /* bits */; + // Number of bytes for load B alignment for TMA load. + auto const numBytesAlignmentLoadB = 1024; + // No need to reuse the first chunk. + auto const reuseChunksSmemLoadB = false; + // Add info. + smemChunkNames.emplace_back("smemLoadB"); + numBytesAndAlignmentPerSmemChunk.emplace_back( + std::make_pair(numSmemBytesLoadB, numBytesAlignmentLoadB)); + firstChunkReuseSmem.emplace_back(reuseChunksSmemLoadB); + } + + // SmemBShuffle + // FIXME: we should be able either: + // - Do modification in-place. For that we need to resolve pipeline dependency between + // smemB -> shuffleSmemB -> mma + // - Do 4 TMA SW32 loads or several LDGSTS loads. + { + // Number of bytes in save shuffled B in shared memory. + auto const numSmemBytesLoadB = + numSlicesForSliceK > 1 + ? numStages * tileN * tileK * getNumSmemBitsPerElt(dtypeB, mMmaKind) / 8 /* bits */ + : 0; + // Number of bytes for load B alignment for TMA load. + auto const numBytesAlignmentLoadB = 1024; + // No need to reuse the first chunk. + auto const reuseChunksSmemLoadB = false; + + // Add info. + smemChunkNames.emplace_back("smemBShuffle"); + numBytesAndAlignmentPerSmemChunk.emplace_back( + std::make_pair(numSmemBytesLoadB, numBytesAlignmentLoadB)); + firstChunkReuseSmem.emplace_back(reuseChunksSmemLoadB); + } + + // GmemC + // FIXME we might need to fix this for GemmGatedAct, it needs less SMEM to store gated output. + for (int resIdx = 0; resIdx < 2; ++resIdx) { + // Type of the data in the SMEM for GmemC + auto dtypeSmemC = dtypeC; + if (allReduceAlgo == AllReduceAlgo::TwoShot || numSlicesForSplitK > 1) { + dtypeSmemC = dtypeAcc; + } + // Smem is used for GmemC output tile for TMA store and SplitK in CGA. + bool usesSmemForGmemC = useTmaStore || doesSplitKUseDsmem(splitK); + // SMEM for at leader CTA in DSMEM split-k contains K slices. + auto extraGmemCMultiplier = doesSplitKUseDsmem(splitK) ? numSlicesForSplitK : 1; + if (numSlicesForSliceK > 1) { + // TileN is expanded in N dimension for slice-K. + extraGmemCMultiplier *= numSlicesForSliceK; + } + + if (resIdx != 0 && !useDeepSeekFp8) { + // No data for Epilogue1 in case of non-DeepSeek GEMM. + extraGmemCMultiplier = 0; + } + + // Number of bytes to store the output in smem. + auto const numBytesSmemStoreC = usesSmemForGmemC + ? extraGmemCMultiplier * epilogueTileM * epilogueTileN * + tg::dtypeGetNumBits(dtypeSmemC) / 8 /* bits */ + : 0; + // Number of bytes for store C alignment for TMA store. + auto const numBytesAlignmentStoreC = 1024; + // gmemC reuses loadAb memory for split-K in DSMEM. + // Epilogue1 does not reuse and continues after the memory allocated Epilogue0 + // NOTE: we can always reuse loadAb SMEM as long as we don't have persistent scheduler. + auto const reuseFirstChunksSmemStoreC = + doesSplitKUseDsmem(splitK) && resIdx == 0 && !usePersistentScheduler; + + // Add info. + smemChunkNames.emplace_back("smemGmemC" + std::to_string(resIdx)); + numBytesAndAlignmentPerSmemChunk.emplace_back( + std::make_pair(numBytesSmemStoreC, numBytesAlignmentStoreC)); + firstChunkReuseSmem.emplace_back(reuseFirstChunksSmemStoreC); + } + + // RowMax + { + // Number of dqSfsC per CTA. + auto const numDqSfsCPerCta = transposeMmaOutput ? tileM : tileN; + // Number of bytes for rowMax in SMEM. + auto const numBytesSmemRowMax = (useDeepSeekFp8 ? numDqSfsCPerCta : 0) * + tg::dtypeGetNumBits(tg::Dtype::Fp32) / 8 /* bits */; + // Number of bytes alignment for rowMax in SMEM. + auto const numBytesAlignmentRowMax = 16; + + // Add info. + smemChunkNames.emplace_back("smemRowMax"); + numBytesAndAlignmentPerSmemChunk.emplace_back( + std::make_pair(numBytesSmemRowMax, numBytesAlignmentRowMax)); + firstChunkReuseSmem.emplace_back(false); + } + + // SliceK + { + // Real tile size before slice-K reduction. + auto const tileSize = + numSlicesForSliceK > 1 ? numSlicesForSliceK * tileM * numSlicesForSliceK * tileN : 0; + // Number of bytes for tile in SMEM. + auto const numBytesSmemTile = tileSize * tg::dtypeGetNumBits(dtypeAcc) / 8 /* bits */; + // Number of bytes alignment for rowMax in SMEM. + auto const numBytesAlignmentTile = 16; + + // Add info. + smemChunkNames.emplace_back("smemSliceK"); + numBytesAndAlignmentPerSmemChunk.emplace_back( + std::make_pair(numBytesSmemTile, numBytesAlignmentTile)); + firstChunkReuseSmem.emplace_back(false); + } + + // Per-token Scale Factors + { + // Number of bytes for per-token scale factors + auto const numBytesSmemPerTokenSf = (usePerTokenSfA ? (tileM) * sizeof(float) : 0) + + (usePerTokenSfB ? (tileN) * sizeof(float) : 0); + // Number of bytes alignment for per-token scale factors + auto const numBytesAlignmentPerTokenSf = 16; + // Add info. + smemChunkNames.emplace_back("smemPerTokenSf"); + numBytesAndAlignmentPerSmemChunk.emplace_back( + std::make_pair(numBytesSmemPerTokenSf, numBytesAlignmentPerTokenSf)); + firstChunkReuseSmem.emplace_back(false); + } + + // Bias + { + int32_t numBytesSmemBias = 0; + if (isBiasTypeN(biasType)) { + numBytesSmemBias = tileN * sizeof(float); + } else if (isBiasTypeM(biasType)) { + numBytesSmemBias = tileM * sizeof(float); + } else if (isBiasTypeMn(biasType)) { + numBytesSmemBias = tileM * tileN * sizeof(float); + } + // Number of bytes alignment for bias + auto const numBytesAlignmentBias = 16; + // Add info. + smemChunkNames.emplace_back("smemBias"); + numBytesAndAlignmentPerSmemChunk.emplace_back( + std::make_pair(numBytesSmemBias, numBytesAlignmentBias)); + firstChunkReuseSmem.emplace_back(false); + } + + // Per-block absolute maximum for multi-warp reduction. + { + // Number of bytes: number of epilogue warps * number of tile columns. + // TODO: avoid allocating this memory when it's not needed (it's only for MxFp8 + fusedAct) + auto const numBytesSmemBlockAmax = transposeMmaOutput ? 4 * tileN * sizeof(float) : 0; + // Number of bytes alignment. + auto const numBytesAlignmentBlockAmax = 16; + // Add info. + smemChunkNames.emplace_back("smemBlockAmax"); + numBytesAndAlignmentPerSmemChunk.emplace_back( + std::make_pair(numBytesSmemBlockAmax, numBytesAlignmentBlockAmax)); + firstChunkReuseSmem.emplace_back(false); + } + + // SmemConstSfBuf + // A buffer used to copy constant values to TMEM. + { + // Do we need the buffer? + bool const useConstSfBuf = dtypeB == tg::Dtype::E4m3 && dtypeMmaB == tg::Dtype::MxE4m3; + // Number of bytes for the buffer. + auto const numSmemBytesConstSfBuf = useConstSfBuf ? 512 : 0; + // Number of bytes for the alignment of the buffer. + auto const numBytesAlignmentConstSfBuf = 16; + // No need to reuse the first chunk. + auto const reuseChunksSmemConstSfBuf = false; + + // Add info. + smemChunkNames.emplace_back("smemConstSfBuf"); + numBytesAndAlignmentPerSmemChunk.emplace_back( + std::make_pair(numSmemBytesConstSfBuf, numBytesAlignmentConstSfBuf)); + firstChunkReuseSmem.emplace_back(reuseChunksSmemConstSfBuf); + } + + // Create SMEM helper object. + mSmemAllocatorHelper = + MemAllocatorHelper(numBytesAndAlignmentPerSmemChunk, firstChunkReuseSmem, smemChunkNames); +#if 0 + // E.g., + // Chunk 0 smemLoadA: 32768 bytes, 1024 alignment, false, offset 0 + // Chunk 1 smemLoadB: 32768 bytes, 1024 alignment, false, offset 32768 + // Chunk 2 smemBShuffle: 0 bytes, 1024 alignment, false, offset 65536 + // Chunk 3 smemGmemC0: 65536 bytes, 1024 alignment, true, offset 0 + // Chunk 4 smemGmemC1: 65536 bytes, 1024 alignment, false, offset 65536 + // Chunk 5 smemRowMax: 512 bytes, 16 alignment, false, offset 131072 + // Chunk 6 smemSliceK: 0 bytes, 16 alignment, false, offset 131584 + // Chunk 7 smemPerTokenSf: 0 bytes, 16 alignment, false, offset 131584 + mSmemAllocatorHelper.print(); +#endif + } + + // + // TMEM + // + // [..D..][..A..][.SfA.][.SfB.] + { + std::vector> numBytesAndAlignmentPerTmemChunk; + std::vector firstChunkReuseTmem; + std::vector tmemChunkNames; + // Matrix D + { + // Number of columns for accumulators. + auto const numTmemColsD = numSlicesForSliceK * tileN * numStagesMma * + tg::dtypeGetNumBits(dtypeAcc) / + tg::dtypeGetNumBits(tg::Dtype::UInt32); + // Number of columns for D alignment. + auto const numColsAlignmentD = 2; + // No need to reuse TMEM. + auto const reuseChunksTmemD = false; + + // Add info. + tmemChunkNames.emplace_back("tmemD"); + numBytesAndAlignmentPerTmemChunk.emplace_back( + std::make_pair(numTmemColsD, numColsAlignmentD)); + firstChunkReuseTmem.emplace_back(reuseChunksTmemD); + } + + // Matrix A + { + // We use TMEM for A if we use slice-K or if we need to cast A. + bool const useTmemA = (numSlicesForSliceK > 1) || (dtypeMmaA != dtypeA); + // Number of columns for A. + auto const numTmemColsA = + useTmemA ? numStages * tileK / + (numSlicesForSliceK * tg::dtypeGetNumBits(tg::Dtype::UInt32) / + tg::dtypeGetNumBits(dtypeMmaA)) + : 0; + // Number of columns for A alignment. + auto const numColsAlignmentA = 4; + // No need to reuse TMEM. + auto const reuseChunksTmemA = false; + + // Add info. + tmemChunkNames.emplace_back("tmemA"); + numBytesAndAlignmentPerTmemChunk.emplace_back( + std::make_pair(numTmemColsA, numColsAlignmentA)); + firstChunkReuseTmem.emplace_back(reuseChunksTmemA); + } + + // Sf A + { + // Does the MMA require block scales in TMEM for A? + bool const useBlockScalingA = tg::dtypeIsBlockFmt(dtypeMmaA); + // Are the block scales constant? + bool const useConstSfA = useBlockScalingA && !tg::dtypeIsBlockFmt(dtypeA); + // Number of columns for scaling factors of A. + auto const numTmemColsSfA = + useConstSfA + ? tg::roundUp((tileK / 64) * 2 * tg::ceilDiv(tileM, 64), 4) + : (useBlockScalingA ? ((tileK / 64) * 2 * tg::ceilDiv(tileM, 64)) * numStages : 0); + // Number of columns for Sf alignment. + auto const numColsAlignmentSfA = 4; + // No need to reuse TMEM. + auto const reuseChunksTmemSfA = false; + + // Add info. + tmemChunkNames.emplace_back("tmemSfA"); + numBytesAndAlignmentPerTmemChunk.emplace_back( + std::make_pair(numTmemColsSfA, numColsAlignmentSfA)); + firstChunkReuseTmem.emplace_back(reuseChunksTmemSfA); + } + + // Sf B + { + // Does the MMA require block scales in TMEM for B? + bool const useBlockScalingB = tg::dtypeIsBlockFmt(dtypeMmaB); + // Are the block scales constant? + bool const useConstSfB = useBlockScalingB && !tg::dtypeIsBlockFmt(dtypeB); + // Number of columns for scaling factors of B. + auto const numTmemColsSfB = + useConstSfB + ? tg::roundUp((tileK / 64) * 2 * tg::ceilDiv(tileN, 64), 4) + : (useBlockScalingB ? ((tileK / 64) * 2 * tg::ceilDiv(tileN, 64)) * numStages : 0); + // Number of columns for Sf alignment. + auto const numColsAlignmentSfB = 4; + // No need to reuse TMEM. + auto const reuseChunksTmemSfB = false; + + // Add info. + tmemChunkNames.emplace_back("tmemSfB"); + numBytesAndAlignmentPerTmemChunk.emplace_back( + std::make_pair(numTmemColsSfB, numColsAlignmentSfB)); + firstChunkReuseTmem.emplace_back(reuseChunksTmemSfB); + } + + // Create TMEM helper object. + mTmemAllocatorHelper = + MemAllocatorHelper(numBytesAndAlignmentPerTmemChunk, firstChunkReuseTmem, tmemChunkNames); + } + } + + public: + // The MMA kind. + tg::MmaKind mMmaKind; + // Helper for SMEM allocation. + MemAllocatorHelper mSmemAllocatorHelper; + // Helper for TMEM allocation. + MemAllocatorHelper mTmemAllocatorHelper; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline int32_t getSmemBufferSize(KernelTraits traits) { + return traits.mSmemAllocatorHelper.getTotalSize(); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline int32_t getTmemBufferSize(KernelTraits traits) { + return traits.mTmemAllocatorHelper.getTotalSize(); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// +// +// Starting address of each SMEM buffer. +// +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline int32_t getSmemOffsetLoadA(KernelTraits traits) { + return traits.mSmemAllocatorHelper.getChunkOffset(0); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline int32_t getSmemOffsetLoadB(KernelTraits traits) { + return traits.mSmemAllocatorHelper.getChunkOffset(1); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline int32_t getSmemOffsetLoadAb(KernelTraits traits) { return getSmemOffsetLoadA(traits); } + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline int32_t getSmemOffsetLoadShuffleB(KernelTraits traits) { + return traits.mSmemAllocatorHelper.getChunkOffset(2); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline int32_t getSmemOffsetGmemC(KernelTraits traits, int resIdx = 0) { + return traits.mSmemAllocatorHelper.getChunkOffset(3 + resIdx); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline int32_t getSmemOffsetRowMax(KernelTraits traits) { + return traits.mSmemAllocatorHelper.getChunkOffset(5); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline int32_t getSmemOffsetSliceK(KernelTraits traits) { + return traits.mSmemAllocatorHelper.getChunkOffset(6); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline int32_t getSmemOffsetPerTokenSf(KernelTraits traits) { + return traits.mSmemAllocatorHelper.getChunkOffset(7); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline int32_t getSmemOffsetBias(KernelTraits traits) { + return traits.mSmemAllocatorHelper.getChunkOffset(8); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline int32_t getSmemOffsetBlockAmax(KernelTraits traits) { + return traits.mSmemAllocatorHelper.getChunkOffset(9); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline int32_t getSmemOffsetConstSfBuf(KernelTraits traits) { + return traits.mSmemAllocatorHelper.getChunkOffset(10); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline int32_t isSmemAbRepurposedToGmemC(KernelTraits traits, int resIdx = 0) { + // Be conscious that the index (3 + resIdx) should match the index in getSmemOffsetGmemC(). + return traits.mSmemAllocatorHelper.getFirstChunkReuseFlag(3 + resIdx); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// +// +// Starting address of each TMEM buffer. +// +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline int32_t getTmemOffsetD(KernelTraits traits) { + return traits.mTmemAllocatorHelper.getChunkOffset(0); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline int32_t getTmemOffsetA(KernelTraits traits) { + return traits.mTmemAllocatorHelper.getChunkOffset(1); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline int32_t getTmemOffsetSfA(KernelTraits traits) { + return traits.mTmemAllocatorHelper.getChunkOffset(2); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline int32_t getTmemOffsetSfB(KernelTraits traits) { + return traits.mTmemAllocatorHelper.getChunkOffset(3); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace gemm + +} // namespace gemm diff --git a/include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/TmaDescriptor.h b/include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/TmaDescriptor.h new file mode 100644 index 0000000000000000000000000000000000000000..4560f89ec632be184a184349bb5d48bf2b368f85 --- /dev/null +++ b/include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/TmaDescriptor.h @@ -0,0 +1,293 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & + * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include + +#include "trtllm/gen/DtypeDecl.h" +#include "trtllm/gen/MmaDecl.h" + +#ifdef TLLM_ENABLE_CUDA +#include +#include +#include +#endif + +namespace gemm { + +namespace gemm { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace tg = trtllm::gen; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +#ifdef TLLM_ENABLE_CUDA + +inline CUtensorMap buildNdTmaDescriptor(tg::Dtype dtype, tg::MmaKind mmaKind, + std::vector const& shapes, + std::vector const& strides, + std::vector const& tileShapes, void* gmemAddr, + bool doSwizzle = true) { + // The multiplication factor of the data padding in SMEM. + int32_t padMultiplier = 1; + CUtensorMap desc{}; + // The data type. + CUtensorMapDataType tmaDataFormat{CU_TENSOR_MAP_DATA_TYPE_FLOAT32}; + if (dtype == tg::Dtype::E4m3 || dtype == tg::Dtype::MxE4m3 || dtype == tg::Dtype::UE8m0) { + tmaDataFormat = CU_TENSOR_MAP_DATA_TYPE_UINT8; + } else if (dtype == tg::Dtype::Fp16) { + tmaDataFormat = CU_TENSOR_MAP_DATA_TYPE_FLOAT16; + } else if (dtype == tg::Dtype::Bfloat16) { + tmaDataFormat = CU_TENSOR_MAP_DATA_TYPE_BFLOAT16; + } else if (dtype == tg::Dtype::E2m1) { + tmaDataFormat = CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN8B; + } else if (dtype == tg::Dtype::MxE2m1) { + if (mmaKind == tg::MmaKind::MxFp8Fp6Fp4) { + padMultiplier = 2; + tmaDataFormat = CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN16B; + } else { + // Note: this is used with the MMA kind MxFp4NvFp4 and also when casting to a higher-precision + // type such as Bfloat16 before the MMA. + tmaDataFormat = CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN8B; + } + } else if (dtype == tg::Dtype::Fp32) { + tmaDataFormat = CU_TENSOR_MAP_DATA_TYPE_FLOAT32; + } else { + std::cerr << "buildNdTmaDescriptor: unexpected dtype " << tg::dtypeToString(dtype) << std::endl; + assert(false); + } + + // The swizzle type. + CUtensorMapSwizzle swizzleType{CU_TENSOR_MAP_SWIZZLE_NONE}; + int32_t fastestDimTileSizeBytes = + (tileShapes[0] * tg::dtypeGetNumBits(dtype) * padMultiplier) / /* bits */ 8; + if (doSwizzle) { + if ((fastestDimTileSizeBytes % 128) == 0) { + swizzleType = CU_TENSOR_MAP_SWIZZLE_128B; + } else if ((fastestDimTileSizeBytes % 64) == 0) { + swizzleType = CU_TENSOR_MAP_SWIZZLE_64B; + } else if ((fastestDimTileSizeBytes % 32) == 0) { + swizzleType = CU_TENSOR_MAP_SWIZZLE_32B; + // This path is only for the scaling factors. + } else if ((fastestDimTileSizeBytes % 16) == 0 && + (dtype == tg::Dtype::UE8m0 || dtype == tg::Dtype::E4m3)) { + swizzleType = CU_TENSOR_MAP_SWIZZLE_NONE; + } else { + std::cerr << "buildNdTmaDescriptor: unexpected fastestDimTileSizeBytes " + << fastestDimTileSizeBytes << std::endl; + assert(false); + } + } + + // Check gmem address must be 16B-aligned + assert((reinterpret_cast(gmemAddr) & 0b1111) == 0); // + + // Check shape must be in range [1, 2^32] + int32_t dim = shapes.size(); + // Expect 2 dimensions for regular gemm, 3 dimensions for batched gemm or blocked layout, and 4 + // dimensions for batched gemm with blocked layout. + assert(dim == 2 || dim == 3 || dim == 4); + // Check shape range. + for (int32_t ii = 0; ii < dim; ++ii) { + assert(shapes[ii] >= (uint64_t(1))); // Size must be min 1 + assert(shapes[ii] <= (uint64_t(1) << 32)); // Size must be max 2^32 + } + + // TMA descriptor does not store the zeroth stride and assumes it is 1. + assert(static_cast(strides.size()) == dim); + assert(strides[0] == 1); + + // Build strides in bytes. + // cuTensorMapEncodeTiled ignores the stride of the first dimension (implicitly 1). + std::vector stridesInBytes(dim - 1); + for (int32_t ii = 0; ii < dim - 1; ++ii) { + stridesInBytes[ii] = (strides[ii + 1] * tg::dtypeGetNumBits(dtype)) / /* bits */ 8; + } + + // Set the number of elements in the packed uint32_t element. + auto const numEltsPerUInt32 = 4 * /* bits */ 8 / (tg::dtypeGetNumBits(dtype) * padMultiplier); + // The number of elements in 128B. + auto const numEltsIn128B = numEltsPerUInt32 /*4B*/ * 32; + // The number of tile K hidden size (per token) in each block of shared memory. + auto const numEltsInClampedFastestTileSize = std::min(numEltsIn128B, tileShapes[0]); + + // Build box dim array. If tileShapes is smaller than dim, just fill with 1s. + assert(static_cast(tileShapes.size()) <= dim); + std::vector boxDim(dim, 1); + boxDim[0] = numEltsInClampedFastestTileSize; + for (size_t ii = 1; ii < tileShapes.size(); ++ii) { + if (tileShapes[ii] > 256) { + std::cerr << "buildNdTmaDescriptor: boxDim too large " << tileShapes[ii] << std::endl; + assert(false); + } else { + boxDim[ii] = tileShapes[ii]; + } + } + + // Set tile strides to 1; + std::vector tileStrides(dim, 1); + + // Build the descriptor. + CUresult result = + cuTensorMapEncodeTiled(&desc, tmaDataFormat, + /*tensorRank=*/dim, gmemAddr, shapes.data(), stridesInBytes.data(), + boxDim.data(), tileStrides.data(), + /*interleave=*/CU_TENSOR_MAP_INTERLEAVE_NONE, swizzleType, + /*l2Promotion=*/CU_TENSOR_MAP_L2_PROMOTION_L2_128B, + /*oobFill=*/CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + + if (result != CUDA_SUCCESS) { + std::stringstream ss; + ss << "Error: Failed to initialize the TMA descriptor " << result << std::endl; + + ss << "tmaFormat: " << static_cast(tmaDataFormat) << " dim: " << dim + << " gmem: " << gmemAddr << std::endl; + + ss << "Shape: "; + for (int ii = 0; ii < dim; ++ii) { + ss << shapes[ii] << " "; + } + ss << std::endl; + + ss << "Stride: "; + for (int ii = 0; ii < dim - 1; ++ii) { + ss << stridesInBytes[ii] << " "; + } + ss << std::endl; + + ss << "tileShapes: "; + for (int ii = 0; ii < dim; ++ii) { + ss << boxDim[ii] << " "; + } + ss << std::endl; + + ss << "tileStrides: "; + for (int ii = 0; ii < dim; ++ii) { + ss << tileStrides[ii] << " "; + } + ss << std::endl; + ss << "swizzleType: " << int(swizzleType) << std::endl; + ss << "(in " << __FILE__ << ":" << __LINE__ << ")" << std::endl; + throw std::runtime_error(ss.str()); + } + + return desc; +} + +// TODO: make it work with the above descriptor? +inline CUtensorMap buildSfTmaDescriptor(tg::Dtype dtype, std::vector const& shapes, + std::vector const& strides, + const std::vector& tileShapes, void* gmemAddr) { + CUtensorMap desc{}; + CUtensorMapDataType tmaDataFormat; + if (dtype == tg::Dtype::E4m3 || dtype == tg::Dtype::UE8m0) { + tmaDataFormat = CU_TENSOR_MAP_DATA_TYPE_UINT8; + } else { + std::cerr << "buildSfTmaDescriptor: unexpected dtype " << tg::dtypeToString(dtype) << std::endl; + assert(false); + } + + // No swizzle for scaling factors. + CUtensorMapSwizzle swizzleType = CU_TENSOR_MAP_SWIZZLE_NONE; + + // Check gmem address must be 16B-aligned + assert((reinterpret_cast(gmemAddr) & 0b1111) == 0); // + + // Check shape must be in range [1, 2^32] + int32_t dim = shapes.size(); + // Check shape range. + for (int32_t ii = 0; ii < dim; ++ii) { + assert(shapes[ii] >= (uint64_t(1))); // Size must be min 1 + assert(shapes[ii] <= (uint64_t(1) << 32)); // Size must be max 2^32 + } + + // TMA descriptor does not store the zeroth stride and assumes it is 1. + assert(static_cast(strides.size()) == dim); + assert(strides[0] == 1); + + // Build strides in bytes. + // cuTensorMapEncodeTiled ignores the stride of the first dimension (implicitly 1). + std::vector stridesInBytes(dim - 1); + for (int32_t ii = 0; ii < dim - 1; ++ii) { + stridesInBytes[ii] = (strides[ii + 1] * tg::dtypeGetNumBits(dtype)) / /* bits */ 8; + } + + // Set tile strides to 1; + std::vector tileStrides(dim, 1); + + // Build the descriptor. + CUresult result = cuTensorMapEncodeTiled(/*tensorMap=*/&desc, + /*tensorDataType=*/tmaDataFormat, + /*tensorRank=*/dim, + /*globalAddress=*/gmemAddr, + /*globalDim=*/shapes.data(), + /*globalStrides=*/stridesInBytes.data(), + /*boxDim=*/tileShapes.data(), + /*elementStrides=*/tileStrides.data(), + /*interleave=*/CU_TENSOR_MAP_INTERLEAVE_NONE, + /*swizzle=*/swizzleType, + /*l2Promotion=*/CU_TENSOR_MAP_L2_PROMOTION_L2_128B, + /*oobFill=*/CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + + if (result != CUDA_SUCCESS) { + std::stringstream ss; + ss << "Error: Failed to initialize the TMA descriptor for SF " << result << std::endl; + + ss << "tmaFormat: " << static_cast(tmaDataFormat) << " dim: " << dim + << " gmem: " << gmemAddr << std::endl; + + ss << "shape:"; + for (uint32_t shape_i : shapes) { + ss << " " << shape_i; + } + ss << std::endl; + + ss << "stridesInBytes:"; + for (uint32_t stride_i : stridesInBytes) { + ss << " " << stride_i; + } + ss << std::endl; + + ss << "tileShapes:"; + for (uint32_t tileShape_i : tileShapes) { + ss << " " << tileShape_i; + } + ss << std::endl; + + ss << "tileStrides:"; + for (uint32_t tileStride_i : tileStrides) { + ss << " " << tileStride_i; + } + ss << std::endl; + + ss << "swizzleType: " << int(swizzleType) << std::endl; + ss << "(in " << __FILE__ << ":" << __LINE__ << ")" << std::endl; + throw std::runtime_error(ss.str()); + } + + return desc; +} + +#endif // defined TLLM_ENABLE_CUDA + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace gemm + +} // namespace gemm diff --git a/include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/trtllm/gen/CommonUtils.h b/include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/trtllm/gen/CommonUtils.h new file mode 100644 index 0000000000000000000000000000000000000000..27c7ab7193faab8a5af3100d7bb7819c7d0dc459 --- /dev/null +++ b/include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/trtllm/gen/CommonUtils.h @@ -0,0 +1,43 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & + * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +namespace gemm { + +namespace trtllm { +namespace gen { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +inline T ceilDiv(T m, T n) { + return (m + n - T(1)) / n; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +inline T roundUp(T m, T n) { + return ceilDiv(m, n) * n; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace gen +} // namespace trtllm + +} // namespace gemm diff --git a/include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/trtllm/gen/CudaKernelLauncher.h b/include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/trtllm/gen/CudaKernelLauncher.h new file mode 100644 index 0000000000000000000000000000000000000000..92c88dda163f43f50401d791d0ff7c8bd5357482 --- /dev/null +++ b/include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/trtllm/gen/CudaKernelLauncher.h @@ -0,0 +1,92 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & + * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#ifdef TLLM_ENABLE_CUDA +#include +#include + +#include +#include +#endif +namespace gemm { + +namespace trtllm { +namespace gen { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +#ifdef TLLM_ENABLE_CUDA +inline CUresult launchKernel(void* kernelParams, void* cudaStream, int32_t smemSize, + CUfunction kernel, dim3 block3, dim3 grid3, dim3 cluster3, + bool enablesPdl) { + // Make sure we can launch with that much shared memory. + if (smemSize > 48 * 1024) { + CUresult result = + cuFuncSetAttribute(kernel, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, smemSize); + if (result != CUDA_SUCCESS) { + return result; + } + } + + auto clusterDim = cluster3.x * cluster3.y * cluster3.z; + + CUlaunchConfig launchConfig; + launchConfig.blockDimX = block3.x; + launchConfig.blockDimY = block3.y; + launchConfig.blockDimZ = block3.z; + launchConfig.gridDimX = grid3.x; + launchConfig.gridDimY = grid3.y; + launchConfig.gridDimZ = grid3.z; + launchConfig.hStream = reinterpret_cast(cudaStream); + launchConfig.sharedMemBytes = smemSize; + + CUlaunchAttribute launchAttrs[3]; + launchAttrs[0].id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION; + launchAttrs[0].value.clusterDim.x = cluster3.x; + launchAttrs[0].value.clusterDim.y = cluster3.y; + launchAttrs[0].value.clusterDim.z = cluster3.z; + launchAttrs[1].id = CU_LAUNCH_ATTRIBUTE_CLUSTER_SCHEDULING_POLICY_PREFERENCE; + launchAttrs[1].value.clusterSchedulingPolicyPreference = + (clusterDim > 1) ? CU_CLUSTER_SCHEDULING_POLICY_SPREAD : CU_CLUSTER_SCHEDULING_POLICY_DEFAULT; + launchAttrs[2].id = CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION; + launchAttrs[2].value.programmaticStreamSerializationAllowed = enablesPdl; + launchConfig.attrs = launchAttrs; + launchConfig.numAttrs = 3; + + // Add setting for non-portable cluster size. + if (clusterDim > 8) { + CUresult result = + cuFuncSetAttribute(kernel, CU_FUNC_ATTRIBUTE_NON_PORTABLE_CLUSTER_SIZE_ALLOWED, + 1 // Enable non-portable cluster sizes + ); + if (result != CUDA_SUCCESS) { + return result; + } + } + + // Launch the kernel. + return cuLaunchKernelEx(&launchConfig, kernel, &kernelParams, nullptr); +} +#endif + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace gen +} // namespace trtllm + +} // namespace gemm diff --git a/include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/trtllm/gen/DtypeDecl.h b/include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/trtllm/gen/DtypeDecl.h new file mode 100644 index 0000000000000000000000000000000000000000..0e087769f01f6eb1970827376aab87342a155458 --- /dev/null +++ b/include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/trtllm/gen/DtypeDecl.h @@ -0,0 +1,271 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & + * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include +#include +#include +#include +#ifndef TLLM_GEN_EXPORT_INTERFACE +#include "trtllm/gen/MmaDecl.h" +#else +#include "MmaDecl.h" +#endif + +//////////////////////////////////////////////////////////////////////////////////////////////////// +// +// Be careful when modifying this file as it is included by the generated kernels. For example, do +// not add TLLM_CHECK_* constructs in this file. Thanks! +// +//////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace gemm { + +namespace trtllm { +namespace gen { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +enum class Dtype : uint32_t { + +// We use the following encoding for the types: +// +// Byte 0: Identifier for the type (going from 0 to the number of data types - 1, +// Byte 1: Number of bits in the type, +// Byte 2: Bit 0: Is it an integer? 0x1 if true, 0x0 otherwise; +// Bit 4: is it signed? 0x1 if true, 0x0 otherwise. +// Byte 3: Is it a block format? 0x1 if true, 0x0 otherwise. + +#define TLLM_ENCODE_DTYPE(BlockFormatBit, SignedBit, IntegerBit, NumBits, Uid) \ + uint32_t { \ + (BlockFormatBit << 24) | (SignedBit << 20) | (IntegerBit << 16) | (NumBits << 8) | (Uid) \ + } + + // clang-format off + Bfloat16 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 1u, /*int*/ 0u, /*bits*/ 16u, /*uid*/ 0u), + Bool = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 0u, /*int*/ 1u, /*bits*/ 1u, /*uid*/ 1u), + E2m1 = TLLM_ENCODE_DTYPE(/*block*/ 1u, /*signed*/ 1u, /*int*/ 0u, /*bits*/ 4u, /*uid*/ 2u), + E2m3 = TLLM_ENCODE_DTYPE(/*block*/ 1u, /*signed*/ 1u, /*int*/ 0u, /*bits*/ 6u, /*uid*/ 3u), + E3m2 = TLLM_ENCODE_DTYPE(/*block*/ 1u, /*signed*/ 1u, /*int*/ 0u, /*bits*/ 6u, /*uid*/ 4u), + E4m3 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 1u, /*int*/ 0u, /*bits*/ 8u, /*uid*/ 5u), + E5m2 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 1u, /*int*/ 0u, /*bits*/ 8u, /*uid*/ 6u), + Fp16 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 1u, /*int*/ 0u, /*bits*/ 16u, /*uid*/ 7u), + Fp32 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 1u, /*int*/ 0u, /*bits*/ 32u, /*uid*/ 8u), + Int8 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 1u, /*int*/ 1u, /*bits*/ 8u, /*uid*/ 9u), + Int32 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 1u, /*int*/ 1u, /*bits*/ 32u, /*uid*/ 10u), + Int64 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 1u, /*int*/ 1u, /*bits*/ 64u, /*uid*/ 11u), + MxE2m1 = TLLM_ENCODE_DTYPE(/*block*/ 1u, /*signed*/ 1u, /*int*/ 0u, /*bits*/ 4u, /*uid*/ 12u), + MxE4m3 = TLLM_ENCODE_DTYPE(/*block*/ 1u, /*signed*/ 1u, /*int*/ 0u, /*bits*/ 8u, /*uid*/ 13u), + UE8m0 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 0u, /*int*/ 0u, /*bits*/ 8u, /*uid*/ 14u), + UInt8 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 0u, /*int*/ 1u, /*bits*/ 8u, /*uid*/ 15u), + UInt16 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 0u, /*int*/ 1u, /*bits*/ 16u, /*uid*/ 16u), + UInt32 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 0u, /*int*/ 1u, /*bits*/ 32u, /*uid*/ 17u), + UInt64 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 0u, /*int*/ 1u, /*bits*/ 64u, /*uid*/ 18u), + UInt128 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 0u, /*int*/ 1u, /*bits*/ 128u, /*uid*/ 19u), + Void = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 1u, /*int*/ 0u, /*bits*/ 0u, /*uid*/ 20u), +// clang-format on + +#undef TLLM_ENCODE_DTYPE +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// The number of bits in a data type? +inline int dtypeGetNumBits(Dtype dtype) { + constexpr uint32_t kMask = 0xffu << 8; + return static_cast((static_cast(dtype) & kMask) >> 8); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// Does the format use block scaling? +inline bool dtypeIsBlockFmt(Dtype dtype) { + constexpr uint32_t kMask = 0xffu << 24; + return static_cast((static_cast(dtype) & kMask) >> 24); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// Is a given data type a floating-point type? +inline bool dtypeIsFloat(Dtype dtype) { + constexpr uint32_t kMask = 0x1u << 16; + return dtype != Dtype::Void && 0 == (static_cast(dtype) & kMask); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// Is a given data type an 8-bit floating-point type? +inline bool dtypeIsFp8(Dtype dtype) { return dtype == Dtype::E4m3 || dtype == Dtype::E5m2; } + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// Is a given data type an integer type? +inline bool dtypeIsInt(Dtype dtype) { + constexpr uint32_t kMask = 0x1u << 16; + return (dtype != Dtype::Bool) && (0 != (static_cast(dtype) & kMask)); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// Is a given data type signed? +inline bool dtypeIsSigned(Dtype dtype) { + constexpr uint32_t kMask = 0x1u << 20; + return (0 != (static_cast(dtype) & kMask)); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// For logging and error reporting +inline std::string dtypeToString(Dtype dtype) { + switch (dtype) { + case Dtype::Bfloat16: + return "Bfloat16"; + case Dtype::Bool: + return "Bool"; + case Dtype::E2m1: + return "E2m1"; + case Dtype::E2m3: + return "E2m3"; + case Dtype::E3m2: + return "E3m2"; + case Dtype::E4m3: + return "E4m3"; + case Dtype::E5m2: + return "E5m2"; + case Dtype::Fp16: + return "Fp16"; + case Dtype::Fp32: + return "Fp32"; + case Dtype::Int8: + return "Int8"; + case Dtype::Int32: + return "Int32"; + case Dtype::Int64: + return "Int64"; + case Dtype::MxE4m3: + return "MxE4m3"; + case Dtype::MxE2m1: + return "MxE2m1"; + case Dtype::UE8m0: + return "UE8m0"; + case Dtype::UInt8: + return "UInt8"; + case Dtype::UInt16: + return "UInt16"; + case Dtype::UInt32: + return "UInt32"; + case Dtype::UInt64: + return "UInt64"; + case Dtype::UInt128: + return "UInt128"; + case Dtype::Void: + return "Void"; + default: + assert(false); + return "Unsupported type"; + } +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline Dtype dtypeEltType(Dtype dtype) { + switch (dtype) { + case Dtype::MxE2m1: + return Dtype::E2m1; + case Dtype::MxE4m3: + return Dtype::E4m3; + default: + return dtype; + } +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline int dtypeNumEltsPerSf(Dtype dtype) { + switch (dtype) { + case Dtype::E2m1: + return 16; + case Dtype::MxE2m1: + case Dtype::MxE4m3: + return 32; + default: + assert(false); + return -1; + } +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// Returns the dtype of scaling factors, if applicable. +inline Dtype dtypeGetBlockSfType(Dtype dtype) { + switch (dtype) { + case Dtype::E2m1: + return Dtype::E4m3; + case Dtype::MxE2m1: + case Dtype::MxE4m3: + return Dtype::UE8m0; + default: + assert(false); + return Dtype::Void; + } +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline MmaKind dtypeGetMmaKind(Dtype dtypeA, Dtype dtypeB) { + auto dtypeEltA = dtypeEltType(dtypeA); + auto dtypeEltB = dtypeEltType(dtypeB); + + // Note: the order of the conditions is important here. + if ((dtypeA == Dtype::Fp16 && dtypeB == Dtype::Fp16) || + (dtypeA == Dtype::Bfloat16 && dtypeB == Dtype::Bfloat16)) { + return MmaKind::Fp16; + } + + if ((dtypeA == Dtype::Int8 || dtypeA == Dtype::UInt8) && + (dtypeB == Dtype::Int8 || dtypeB == Dtype::UInt8)) { + return MmaKind::Int8; + } + + // This statement captures both MxE2m1 and E2m1. + if (dtypeEltA == Dtype::E2m1 && dtypeEltB == Dtype::E2m1) { + return MmaKind::MxFp4NvFp4; + } + + if ((dtypeA == Dtype::E4m3 || dtypeA == Dtype::E5m2 || dtypeA == Dtype::E2m3 || + dtypeA == Dtype::E3m2 || dtypeA == Dtype::E2m1) && + (dtypeB == Dtype::E4m3 || dtypeB == Dtype::E5m2 || dtypeB == Dtype::E2m3 || + dtypeB == Dtype::E3m2 || dtypeB == Dtype::E2m1)) { + return MmaKind::Fp8Fp6Fp4; + } + + // At this point we know that both dtypes are Mx types and not both MxE2m1 at the same time. + if ((dtypeEltA == Dtype::E4m3 || dtypeEltA == Dtype::E5m2 || dtypeEltA == Dtype::E2m3 || + dtypeEltA == Dtype::E3m2 || dtypeEltA == Dtype::E2m1) && + (dtypeEltB == Dtype::E4m3 || dtypeEltB == Dtype::E5m2 || dtypeEltB == Dtype::E2m3 || + dtypeEltB == Dtype::E3m2 || dtypeEltB == Dtype::E2m1)) { + return MmaKind::MxFp8Fp6Fp4; + } + return MmaKind::Tf32; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace gen +} // namespace trtllm + +} // namespace gemm diff --git a/include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/trtllm/gen/MmaDecl.h b/include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/trtllm/gen/MmaDecl.h new file mode 100644 index 0000000000000000000000000000000000000000..d56941efbd553703afa2d2b2babc230fc1d248c7 --- /dev/null +++ b/include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/trtllm/gen/MmaDecl.h @@ -0,0 +1,92 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & + * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +namespace gemm { + +namespace trtllm { +namespace gen { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// The kind of the MMA instruction +enum class MmaKind : uint32_t { + // For Blackwell this follows the PTX ISA description of the MMA instructions. + // https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-kind-shapes + + // The MMA type is auto-detected from the dtypes of the input tensors + Auto = 0, + // Supports dtypeA = dtypeB = Fp16 and dtypeD = [Fp16, Fp32] + // or dtypeA = dtypeB = Bfloat16 and dtypeD = [Fp32] + // Corresponds to the kind::f16 of tcgen05.mma. + Fp16 = 1, + // Supports dtypeA/B = [E4m3, E5m2, E2m3, E3m2, E2m1] and dtypeD = [Fp16, Fp32] + // Corresponds to the kind::f8f6f4 of tcgen05.mma. + Fp8Fp6Fp4 = 2, + // Supports dtypeA = dtypeB = [Int8, Uint8] and dtypeD = [Int32] + // Corresponds to the kind::i8 of tcgen05.mma. + Int8 = 3, + // Supports dtypeA = dtypeB = [MxE2m1, E2m1] with block scale [UM8e0, UEm4e3] + // and dtypeD = [Fp32] + // Corresponds to the kind::mxf4nvf4 of tcgen05.mma. + MxFp4NvFp4 = 4, + // Supports dtype dtypeA = dtypeB = [MxE4m3, MxE2m1] with block scale [UM8e0] + // and dtypeD = [Fp32] + // Corresponds to the kind::mxf8f6f4 of tcgen05.mma. + MxFp8Fp6Fp4 = 5, + // Supports dtypeA = dtypeB = Tf32 with dtypeD = [Fp32] + // Corresponds to the kind::tf32 of tcgen05.mma. + Tf32 = 6 +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline bool mmaKindIsBlockFmt(MmaKind mmaKind) { + return mmaKind == MmaKind::MxFp8Fp6Fp4 || mmaKind == MmaKind::MxFp4NvFp4; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// For logging and error reporting +inline std::string mmaKindToString(MmaKind mmaKind) { + switch (mmaKind) { + case MmaKind::Auto: + return "Auto"; + case MmaKind::Fp16: + return "Fp16"; + case MmaKind::Fp8Fp6Fp4: + return "Fp8Fp6Fp4"; + case MmaKind::Int8: + return "Int8"; + case MmaKind::MxFp4NvFp4: + return "MxFp4NvFp4"; + case MmaKind::MxFp8Fp6Fp4: + return "MxFp8Fp6Fp4"; + case MmaKind::Tf32: + return "Tf32"; + default: + assert(false); + return "Unsupported type"; + } +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace gen +} // namespace trtllm + +} // namespace gemm diff --git a/include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/trtllm/gen/SfLayoutDecl.h b/include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/trtllm/gen/SfLayoutDecl.h new file mode 100644 index 0000000000000000000000000000000000000000..25adab890214573d77cea39cb10e7141a0f5e595 --- /dev/null +++ b/include/flashinfer/trtllm/gemm/trtllmGen_gemm_export/trtllm/gen/SfLayoutDecl.h @@ -0,0 +1,96 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & + * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include +#include + +//////////////////////////////////////////////////////////////////////////////////////////////////// +// +// Be careful when modifying this file as it is included by the generated kernels. For example, do +// not add TLLM_CHECK_* constructs in this file. Thanks! +// +//////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace gemm { + +namespace trtllm { +namespace gen { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// This enumeration defines layouts for storing scale factors for FP4, FP6, and FP8 formats. +enum class SfLayout { + // Scale factors are stored in the same order as the associated matrix. + // I.e., the SF buffer is a tensor [m, ⌈n/b⌉], where m, n, and b are respectively the number of + // rows, columns and the block size. + // The SF for the element (i, j) is stored at (i, j/b). + Linear = 0, + + // A tile of 8x4 is stored contiguously. The order of elements inside the tile, and the order + // of tiles, are both row-major. + // I.e., the SF buffer is a tensor [⌈m/8⌉, ⌈n/b/4⌉, 8, 4]. + // The SF for the element (i, j) is stored at (i/8, j/b/4, i%8, (j/b)%4). + R8c4, + + // A tile of 8x16 is stored contiguously. The order of elements inside the tile, and the order + // of tiles, are both row-major. + // I.e., the SF buffer is a tensor [⌈m/8⌉, ⌈n/b/16⌉, 8, 16]. + // The SF for the element (i, j) is stored at (i/8, j/b/16, i%8, (j/b)%16). + // + // NOTE: This is a niche format that is currently used for the weights of the + // LowLatency FP4 kernels. It is not meant as an interchange format. In + // addition to the above requirements it requires n to be a multiple of 256. + R8c16, + + // A tile of 128x4 is stored contiguously. Rows 0-31, 32-63, 64-95 and 96-127 are interleaved + // as illustrated below: + // | 0,0 | 0,1 | 0,2 | 0,3 | 32,0 | 32,1 | 32,2 | 32,3 | ... | 96,3 | + // | 1,0 | 1,1 | 1,2 | 1,3 | 33,0 | 33,1 | 33,2 | 33,3 | ... | 97,3 | + // | ... | ... | ... | ... | ... | ... | ... | ... | ... | ... | + // | 31,0 | 31,1 | 31,2 | 31,3 | 63,0 | 63,1 | 63,2 | 63,3 | ... | 127,3 | + // See https://nvbugspro.nvidia.com/bug/4165523 + // + // I.e., the SF buffer is a tensor [⌈m/128⌉, ⌈n/b/4⌉, 32, 4, 4] + // The SF for the element (i, j) is stored at (i/128, j/b/4, i%32, (i%128)/32, (j/b)%4). + R128c4, +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +inline std::string sfLayoutToString(SfLayout layout) { + switch (layout) { + case SfLayout::Linear: + return "linear"; + case SfLayout::R8c4: + return "8x4"; + case SfLayout::R8c16: + return "8x16"; + case SfLayout::R128c4: + return "128x4"; + default: + assert(false); + return "Unsupported layout"; + } +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace gen +} // namespace trtllm + +} // namespace gemm diff --git a/include/flashinfer/utils.cuh b/include/flashinfer/utils.cuh new file mode 100644 index 0000000000000000000000000000000000000000..9f21f5b2b52e16bc4094b1dbbf834f037f8e015a --- /dev/null +++ b/include/flashinfer/utils.cuh @@ -0,0 +1,377 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef FLASHINFER_UTILS_CUH_ +#define FLASHINFER_UTILS_CUH_ +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +#include "exception.h" + +#define STR_HELPER(x) #x +#define STR(x) STR_HELPER(x) + +// macro to turn off fp16 qk reduction to reduce binary +#ifndef FLASHINFER_ALWAYS_DISUSE_FP16_QK_REDUCTION +#define FLASHINFER_ALWAYS_DISUSE_FP16_QK_REDUCTION 0 +#endif + +#ifndef NDEBUG +#define FLASHINFER_CUDA_CALL(func, ...) \ + { \ + cudaError_t e = (func); \ + if (e != cudaSuccess) { \ + std::cerr << "CUDA Error: " << cudaGetErrorString(e) << " (" << e << ") " << __FILE__ \ + << ": line " << __LINE__ << " at function " << STR(func) << std::endl; \ + return e; \ + } \ + } +#else +#define FLASHINFER_CUDA_CALL(func, ...) \ + { \ + cudaError_t e = (func); \ + if (e != cudaSuccess) { \ + return e; \ + } \ + } +#endif + +#define DISPATCH_USE_FP16_QK_REDUCTION(use_fp16_qk_reduction, USE_FP16_QK_REDUCTION, ...) \ + if (use_fp16_qk_reduction) { \ + FLASHINFER_ERROR("FP16_QK_REDUCTION disabled at compile time"); \ + } else { \ + constexpr bool USE_FP16_QK_REDUCTION = false; \ + __VA_ARGS__ \ + } + +#define DISPATCH_NUM_MMA_Q(num_mma_q, NUM_MMA_Q, ...) \ + if (num_mma_q == 1) { \ + constexpr size_t NUM_MMA_Q = 1; \ + __VA_ARGS__ \ + } else if (num_mma_q == 2) { \ + constexpr size_t NUM_MMA_Q = 2; \ + __VA_ARGS__ \ + } else { \ + std::ostringstream err_msg; \ + err_msg << "Unsupported num_mma_q: " << num_mma_q; \ + FLASHINFER_ERROR(err_msg.str()); \ + } + +#define DISPATCH_NUM_MMA_KV(max_mma_kv, NUM_MMA_KV, ...) \ + if (max_mma_kv >= 8) { \ + constexpr size_t NUM_MMA_KV = 8; \ + __VA_ARGS__ \ + } else if (max_mma_kv >= 4) { \ + constexpr size_t NUM_MMA_KV = 4; \ + __VA_ARGS__ \ + } else if (max_mma_kv >= 2) { \ + constexpr size_t NUM_MMA_KV = 2; \ + __VA_ARGS__ \ + } else if (max_mma_kv >= 1) { \ + constexpr size_t NUM_MMA_KV = 1; \ + __VA_ARGS__ \ + } else { \ + std::ostringstream err_msg; \ + err_msg << "Unsupported max_mma_kv: " << max_mma_kv; \ + FLASHINFER_ERROR(err_msg.str()); \ + } + +#define DISPATCH_CTA_TILE_Q(cta_tile_q, CTA_TILE_Q, ...) \ + switch (cta_tile_q) { \ + case 128: { \ + constexpr uint32_t CTA_TILE_Q = 128; \ + __VA_ARGS__ \ + break; \ + } \ + case 64: { \ + constexpr uint32_t CTA_TILE_Q = 64; \ + __VA_ARGS__ \ + break; \ + } \ + case 16: { \ + constexpr uint32_t CTA_TILE_Q = 16; \ + __VA_ARGS__ \ + break; \ + } \ + default: { \ + std::ostringstream err_msg; \ + err_msg << "Unsupported cta_tile_q: " << cta_tile_q; \ + FLASHINFER_ERROR(err_msg.str()); \ + } \ + } + +#define DISPATCH_GQA_GROUP_SIZE(group_size, GROUP_SIZE, ...) \ + if (group_size == 1) { \ + constexpr size_t GROUP_SIZE = 1; \ + __VA_ARGS__ \ + } else if (group_size == 2) { \ + constexpr size_t GROUP_SIZE = 2; \ + __VA_ARGS__ \ + } else if (group_size == 3) { \ + constexpr size_t GROUP_SIZE = 3; \ + __VA_ARGS__ \ + } else if (group_size == 4) { \ + constexpr size_t GROUP_SIZE = 4; \ + __VA_ARGS__ \ + } else if (group_size == 8) { \ + constexpr size_t GROUP_SIZE = 8; \ + __VA_ARGS__ \ + } else { \ + std::ostringstream err_msg; \ + err_msg << "Unsupported group_size: " << group_size; \ + FLASHINFER_ERROR(err_msg.str()); \ + } + +#define DISPATCH_MASK_MODE(mask_mode, MASK_MODE, ...) \ + switch (mask_mode) { \ + case MaskMode::kNone: { \ + constexpr MaskMode MASK_MODE = MaskMode::kNone; \ + __VA_ARGS__ \ + break; \ + } \ + case MaskMode::kCausal: { \ + constexpr MaskMode MASK_MODE = MaskMode::kCausal; \ + __VA_ARGS__ \ + break; \ + } \ + case MaskMode::kCustom: { \ + constexpr MaskMode MASK_MODE = MaskMode::kCustom; \ + __VA_ARGS__ \ + break; \ + } \ + case MaskMode::kMultiItemScoring: { \ + constexpr MaskMode MASK_MODE = MaskMode::kMultiItemScoring; \ + __VA_ARGS__ \ + break; \ + } \ + default: { \ + std::ostringstream err_msg; \ + err_msg << "Unsupported mask_mode: " << int(mask_mode); \ + FLASHINFER_ERROR(err_msg.str()); \ + } \ + } + +// convert head_dim to compile-time constant +#define DISPATCH_HEAD_DIM(head_dim, HEAD_DIM, ...) \ + switch (head_dim) { \ + case 64: { \ + constexpr size_t HEAD_DIM = 64; \ + __VA_ARGS__ \ + break; \ + } \ + case 128: { \ + constexpr size_t HEAD_DIM = 128; \ + __VA_ARGS__ \ + break; \ + } \ + case 256: { \ + constexpr size_t HEAD_DIM = 256; \ + __VA_ARGS__ \ + break; \ + } \ + case 512: { \ + constexpr size_t HEAD_DIM = 512; \ + __VA_ARGS__ \ + break; \ + } \ + default: { \ + std::ostringstream err_msg; \ + err_msg << "Unsupported head_dim: " << head_dim; \ + FLASHINFER_ERROR(err_msg.str()); \ + } \ + } + +#define DISPATCH_POS_ENCODING_MODE(pos_encoding_mode, POS_ENCODING_MODE, ...) \ + switch (pos_encoding_mode) { \ + case PosEncodingMode::kNone: { \ + constexpr PosEncodingMode POS_ENCODING_MODE = PosEncodingMode::kNone; \ + __VA_ARGS__ \ + break; \ + } \ + case PosEncodingMode::kRoPELlama: { \ + constexpr PosEncodingMode POS_ENCODING_MODE = PosEncodingMode::kRoPELlama; \ + __VA_ARGS__ \ + break; \ + } \ + case PosEncodingMode::kALiBi: { \ + constexpr PosEncodingMode POS_ENCODING_MODE = PosEncodingMode::kALiBi; \ + __VA_ARGS__ \ + break; \ + } \ + default: { \ + std::ostringstream err_msg; \ + err_msg << "Unsupported pos_encoding_mode: " << int(pos_encoding_mode); \ + FLASHINFER_ERROR(err_msg.str()); \ + } \ + } + +#define DISPATCH_ALIGNED_VEC_SIZE(aligned_vec_size, ALIGNED_VEC_SIZE, ...) \ + switch (aligned_vec_size) { \ + case 16: { \ + constexpr size_t ALIGNED_VEC_SIZE = 16; \ + __VA_ARGS__ \ + break; \ + } \ + case 8: { \ + constexpr size_t ALIGNED_VEC_SIZE = 8; \ + __VA_ARGS__ \ + break; \ + } \ + case 4: { \ + constexpr size_t ALIGNED_VEC_SIZE = 4; \ + __VA_ARGS__ \ + break; \ + } \ + case 2: { \ + constexpr size_t ALIGNED_VEC_SIZE = 2; \ + __VA_ARGS__ \ + break; \ + } \ + case 1: { \ + constexpr size_t ALIGNED_VEC_SIZE = 1; \ + __VA_ARGS__ \ + break; \ + } \ + default: { \ + std::ostringstream err_msg; \ + err_msg << "Unsupported aligned_vec_size: " << aligned_vec_size; \ + FLASHINFER_ERROR(err_msg.str()); \ + } \ + } + +#define DISPATCH_COMPUTE_CAP_DECODE_NUM_STAGES_SMEM(compute_capacity, NUM_STAGES_SMEM, ...) \ + if (compute_capacity.first >= 8) { \ + constexpr uint32_t NUM_STAGES_SMEM = 2; \ + __VA_ARGS__ \ + } else { \ + constexpr uint32_t NUM_STAGES_SMEM = 1; \ + __VA_ARGS__ \ + } + +namespace flashinfer { + +template +__forceinline__ __device__ __host__ T1 ceil_div(const T1 x, const T2 y) { + return (x + y - 1) / y; +} + +template +__forceinline__ __device__ __host__ T1 round_up(const T1 x, const T2 y) { + return ceil_div(x, y) * y; +} + +inline std::pair GetCudaComputeCapability() { + int device_id = 0; + cudaGetDevice(&device_id); + int major = 0, minor = 0; + cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, device_id); + cudaDeviceGetAttribute(&minor, cudaDevAttrComputeCapabilityMinor, device_id); + return std::make_pair(major, minor); +} + +template +inline void DebugPrintCUDAArray(T* device_ptr, size_t size, std::string prefix = "") { + std::vector host_array(size); + std::cout << prefix; + cudaMemcpy(host_array.data(), device_ptr, size * sizeof(T), cudaMemcpyDeviceToHost); + for (size_t i = 0; i < size; ++i) { + std::cout << host_array[i] << " "; + } + std::cout << std::endl; +} + +inline uint32_t FA2DetermineCtaTileQ(int64_t avg_packed_qo_len, uint32_t head_dim) { + if (avg_packed_qo_len > 64 && head_dim < 256) { + return 128; + } else { + auto compute_capacity = GetCudaComputeCapability(); + if (compute_capacity.first >= 8) { + // Ampere or newer + if (avg_packed_qo_len > 16) { + // avg_packed_qo_len <= 64 + return 64; + } else { + // avg_packed_qo_len <= 16 + return 16; + } + } else { + // NOTE(Zihao): not enough shared memory on Turing for 1x4 warp layout + return 64; + } + } +} + +#define LOOP_SPLIT_MASK(iter, COND1, COND2, ...) \ + { \ + _Pragma("unroll 1") for (; (COND1); (iter) -= 1) { \ + constexpr bool WITH_MASK = true; \ + __VA_ARGS__ \ + } \ + _Pragma("unroll 1") for (; (COND2); (iter) -= 1) { \ + constexpr bool WITH_MASK = false; \ + __VA_ARGS__ \ + } \ + } + +/*! + * \brief Return x - y if x > y, otherwise return 0. + */ +__device__ __forceinline__ uint32_t sub_if_greater_or_zero(uint32_t x, uint32_t y) { + return (x > y) ? x - y : 0U; +} + +__device__ __forceinline__ void swap(uint32_t& a, uint32_t& b) { + uint32_t tmp = a; + a = b; + b = tmp; +} + +__device__ __forceinline__ uint32_t dim2_offset(const uint32_t& dim_a, const uint32_t& idx_b, + const uint32_t& idx_a) { + return idx_b * dim_a + idx_a; +} + +__device__ __forceinline__ uint32_t dim3_offset(const uint32_t& dim_b, const uint32_t& dim_a, + const uint32_t& idx_c, const uint32_t& idx_b, + const uint32_t& idx_a) { + return (idx_c * dim_b + idx_b) * dim_a + idx_a; +} + +__device__ __forceinline__ uint32_t dim4_offset(const uint32_t& dim_c, const uint32_t& dim_b, + const uint32_t& dim_a, const uint32_t& idx_d, + const uint32_t& idx_c, const uint32_t& idx_b, + const uint32_t& idx_a) { + return ((idx_d * dim_c + idx_c) * dim_b + idx_b) * dim_a + idx_a; +} + +#define DEFINE_HAS_MEMBER(member) \ + template \ + struct has_##member : std::false_type {}; \ + template \ + struct has_##member().member)>> : std::true_type {}; \ + template \ + inline constexpr bool has_##member##_v = has_##member::value; + +} // namespace flashinfer + +#endif // FLASHINFER_UTILS_CUH_ diff --git a/include/flashinfer/vec_dtypes.cuh b/include/flashinfer/vec_dtypes.cuh new file mode 100644 index 0000000000000000000000000000000000000000..fa403f7b933b94b2d2d09857f4ea13545e2644c6 --- /dev/null +++ b/include/flashinfer/vec_dtypes.cuh @@ -0,0 +1,1858 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef VEC_DTYPES_CUH_ +#define VEC_DTYPES_CUH_ + +#include +#include +#include +#include + +#include + +namespace flashinfer { + +#if (!defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 900)) +#define FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED +#endif + +#define FLASHINFER_INLINE inline __attribute__((always_inline)) __device__ + +__device__ __forceinline__ void st_global_release(int4 const& val, int4* addr) { + asm volatile("st.release.global.sys.v4.b32 [%4], {%0, %1, %2, %3};" ::"r"(val.x), "r"(val.y), + "r"(val.z), "r"(val.w), "l"(addr)); +} + +__device__ __forceinline__ int4 ld_global_acquire(int4* addr) { + int4 val; + asm volatile("ld.acquire.global.sys.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(val.x), "=r"(val.y), "=r"(val.z), "=r"(val.w) + : "l"(addr)); + return val; +} + +__device__ __forceinline__ void st_global_volatile(int4 const& val, int4* addr) { + asm volatile("st.volatile.global.v4.b32 [%4], {%0, %1, %2, %3};" ::"r"(val.x), "r"(val.y), + "r"(val.z), "r"(val.w), "l"(addr)); +} + +__device__ __forceinline__ int4 ld_global_volatile(int4* addr) { + int4 val; + asm volatile("ld.volatile.global.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(val.x), "=r"(val.y), "=r"(val.z), "=r"(val.w) + : "l"(addr)); + return val; +} + +#if (__CUDACC_VER_MAJOR__ * 10000 + __CUDACC_VER_MINOR__ * 100 < 120200) && \ + (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 800)) +// CUDA version < 12.2 and GPU architecture < 80 +FLASHINFER_INLINE __nv_bfloat162 make_bfloat162(const __nv_bfloat16 x, const __nv_bfloat16 y) { + __nv_bfloat162 t; + t.x = x; + t.y = y; + return t; +} + +FLASHINFER_INLINE __nv_bfloat16 __hmul(const __nv_bfloat16 a, const __nv_bfloat16 b) { + __nv_bfloat16 val; + const float fa = __bfloat162float(a); + const float fb = __bfloat162float(b); + // avoid ftz in device code + val = __float2bfloat16(__fmaf_ieee_rn(fa, fb, -0.0f)); + return val; +} + +FLASHINFER_INLINE __nv_bfloat162 __hmul2(const __nv_bfloat162 a, const __nv_bfloat162 b) { + __nv_bfloat162 val; + val.x = __hmul(a.x, b.x); + val.y = __hmul(a.y, b.y); + return val; +} + +FLASHINFER_INLINE __nv_bfloat162 __floats2bfloat162_rn(const float a, const float b) { + __nv_bfloat162 val; + val = __nv_bfloat162(__float2bfloat16_rn(a), __float2bfloat16_rn(b)); + return val; +} + +FLASHINFER_INLINE __nv_bfloat162 __float22bfloat162_rn(const float2 a) { + __nv_bfloat162 val = __floats2bfloat162_rn(a.x, a.y); + return val; +} +FLASHINFER_INLINE float2 __bfloat1622float2(const __nv_bfloat162 a) { + float hi_float; + float lo_float; + lo_float = __internal_bfloat162float(((__nv_bfloat162_raw)a).x); + hi_float = __internal_bfloat162float(((__nv_bfloat162_raw)a).y); + return make_float2(lo_float, hi_float); +} +#endif + +/******************* vec_t type cast *******************/ + +template +struct vec_cast { + template + FLASHINFER_INLINE static void cast(dst_t* dst, const src_t* src) { +#pragma unroll + for (size_t i = 0; i < vec_size; ++i) { + dst[i] = (dst_t)src[i]; + } + } +}; + +template <> +struct vec_cast<__nv_fp8_e4m3, float> { + template + FLASHINFER_INLINE static void cast(__nv_fp8_e4m3* dst, const float* src) { + if constexpr (vec_size == 1) { + dst[0] = __nv_fp8_e4m3(src[0]); + } else { +#pragma unroll + for (size_t i = 0; i < vec_size / 2; ++i) { + ((__nv_fp8x2_storage_t*)dst)[i] = + __nv_cvt_float2_to_fp8x2(((float2*)src)[i], __NV_SATFINITE, __NV_E4M3); + } + } + } +}; + +template <> +struct vec_cast<__nv_fp8_e5m2, float> { + template + FLASHINFER_INLINE static void cast(__nv_fp8_e5m2* dst, const float* src) { + if constexpr (vec_size == 1) { + dst[0] = __nv_fp8_e5m2(src[0]); + } else { +#pragma unroll + for (size_t i = 0; i < vec_size / 2; ++i) { + ((__nv_fp8x2_storage_t*)dst)[i] = + __nv_cvt_float2_to_fp8x2(((float2*)src)[i], __NV_SATFINITE, __NV_E5M2); + } + } + } +}; + +template <> +struct vec_cast { + template + FLASHINFER_INLINE static void cast(float* dst, const half* src) { + if constexpr (vec_size == 1) { + dst[0] = (float)src[0]; + } else { +#pragma unroll + for (size_t i = 0; i < vec_size / 2; ++i) { + ((float2*)dst)[i] = __half22float2(((half2*)src)[i]); + } + } + } +}; + +template <> +struct vec_cast { + template + FLASHINFER_INLINE static void cast(half* dst, const float* src) { + if constexpr (vec_size == 1) { + dst[0] = __float2half(src[0]); + } else { +#pragma unroll + for (size_t i = 0; i < vec_size / 2; ++i) { + ((half2*)dst)[i] = __float22half2_rn(((float2*)src)[i]); + } + } + } +}; + +template +constexpr FLASHINFER_INLINE int get_exponent_bits() { + if constexpr (std::is_same_v) { + return 4; + } else if constexpr (std::is_same_v) { + return 5; + } else if constexpr (std::is_same_v) { + return 5; + } else if constexpr (std::is_same_v) { + return 8; + } +} + +template +constexpr FLASHINFER_INLINE int get_mantissa_bits() { + if constexpr (std::is_same_v) { + return 3; + } else if constexpr (std::is_same_v) { + return 2; + } else if constexpr (std::is_same_v) { + return 11; + } else if constexpr (std::is_same_v) { + return 7; + } +} + +/*! + * \brief Fallback to software fast dequant implementation if hardware dequantization is not + * available. + * \note Inspired by Marlin's fast dequantization, but here we don't have to permute + * weights order. + * \ref + * https://github.com/vllm-project/vllm/blob/6dffa4b0a6120159ef2fe44d695a46817aff65bc/csrc/quantization/fp8/fp8_marlin.cu#L120 + */ +template +__device__ void fast_dequant_f8f16x4(uint32_t* input, uint2* output) { + uint32_t q = *input; + if constexpr (std::is_same_v && std::is_same_v) { + output->x = __byte_perm(0U, q, 0x5140); + output->y = __byte_perm(0U, q, 0x7362); + } else { + constexpr int FP8_EXPONENT = get_exponent_bits(); + constexpr int FP8_MANTISSA = get_mantissa_bits(); + constexpr int FP16_EXPONENT = get_exponent_bits(); + + constexpr int RIGHT_SHIFT = FP16_EXPONENT - FP8_EXPONENT; + // Calculate MASK for extracting mantissa and exponent + constexpr int MASK1 = 0x80000000; + constexpr int MASK2 = MASK1 >> (FP8_EXPONENT + FP8_MANTISSA); + constexpr int MASK3 = MASK2 & 0x7fffffff; + constexpr int MASK = MASK3 | (MASK3 >> 16); + q = __byte_perm(q, q, 0x1302); + + // Extract and shift FP8 values to FP16 format + uint32_t Out1 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT); + uint32_t Out2 = ((q << 8) & 0x80008000) | (((q << 8) & MASK) >> RIGHT_SHIFT); + + constexpr int BIAS_OFFSET = (1 << (FP16_EXPONENT - 1)) - (1 << (FP8_EXPONENT - 1)); + // Construct and apply exponent bias + if constexpr (std::is_same_v) { + const half2 bias_reg = __float2half2_rn(float(1 << BIAS_OFFSET)); + + // Convert to half2 and apply bias + *(half2*)&(output->x) = __hmul2(*reinterpret_cast(&Out1), bias_reg); + *(half2*)&(output->y) = __hmul2(*reinterpret_cast(&Out2), bias_reg); + } else { + constexpr uint32_t BIAS = (BIAS_OFFSET + 127) << 23; + const nv_bfloat162 bias_reg = __float2bfloat162_rn(*reinterpret_cast(&BIAS)); + // Convert to bfloat162 and apply bias + *(nv_bfloat162*)&(output->x) = + __hmul2(*reinterpret_cast(&Out1), bias_reg); + *(nv_bfloat162*)&(output->y) = + __hmul2(*reinterpret_cast(&Out2), bias_reg); + } + } +} + +template <> +struct vec_cast { + template + FLASHINFER_INLINE static void cast(nv_bfloat16* dst, const __nv_fp8_e4m3* src) { + if constexpr (vec_size == 1) { + dst[0] = nv_bfloat16(src[0]); + } else if constexpr (vec_size == 2) { + dst[0] = nv_bfloat16(src[0]); + dst[1] = nv_bfloat16(src[1]); + } else { + static_assert(vec_size % 4 == 0, "vec_size must be a multiple of 4"); +#pragma unroll + for (uint32_t i = 0; i < vec_size / 4; ++i) { + fast_dequant_f8f16x4<__nv_fp8_e4m3, nv_bfloat16>((uint32_t*)&src[i * 4], + (uint2*)&dst[i * 4]); + } + } + } +}; + +template <> +struct vec_cast { + template + FLASHINFER_INLINE static void cast(nv_bfloat16* dst, const __nv_fp8_e5m2* src) { + if constexpr (vec_size == 1) { + dst[0] = nv_bfloat16(src[0]); + } else if constexpr (vec_size == 2) { + dst[0] = nv_bfloat16(src[0]); + dst[1] = nv_bfloat16(src[1]); + } else { + static_assert(vec_size % 4 == 0, "vec_size must be a multiple of 4"); +#pragma unroll + for (uint32_t i = 0; i < vec_size / 4; ++i) { + fast_dequant_f8f16x4<__nv_fp8_e5m2, nv_bfloat16>((uint32_t*)&src[i * 4], + (uint2*)&dst[i * 4]); + } + } + } +}; + +template <> +struct vec_cast<__nv_fp8_e4m3, half> { + template + FLASHINFER_INLINE static void cast(__nv_fp8_e4m3* dst, const half* src) { +#ifdef FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED + if constexpr (vec_size == 1) { + dst[0] = __nv_fp8_e4m3(src[0]); + } else { +#pragma unroll + for (size_t i = 0; i < vec_size / 2; ++i) { + uint16_t y; + uint32_t x = *(uint32_t*)&src[i * 2]; + asm volatile("cvt.rn.satfinite.e4m3x2.f16x2 %0, %1;" : "=h"(y) : "r"(x)); + *(uint16_t*)&dst[i * 2] = y; + } + } +#else +#pragma unroll + for (size_t i = 0; i < vec_size; ++i) { + dst[i] = __nv_fp8_e4m3(src[i]); + } +#endif // FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED + } +}; + +template <> +struct vec_cast<__nv_fp8_e5m2, half> { + template + FLASHINFER_INLINE static void cast(__nv_fp8_e5m2* dst, const half* src) { +#ifdef FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED + if constexpr (vec_size == 1) { + dst[0] = __nv_fp8_e5m2(src[0]); + } else { +#pragma unroll + for (size_t i = 0; i < vec_size / 2; ++i) { + uint16_t y; + uint32_t x = *(uint32_t*)&src[i * 2]; + asm volatile("cvt.rn.satfinite.e5m2x2.f16x2 %0, %1;" : "=h"(y) : "r"(x)); + *(uint16_t*)&dst[i * 2] = y; + } + } +#else +#pragma unroll + for (size_t i = 0; i < vec_size; ++i) { + dst[i] = __nv_fp8_e5m2(src[i]); + } +#endif // FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED + } +}; + +template <> +struct vec_cast { + template + FLASHINFER_INLINE static void cast(half* dst, const __nv_fp8_e4m3* src) { +#ifdef FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED + if constexpr (vec_size == 1) { + dst[0] = half(src[0]); + } else { +#pragma unroll + for (size_t i = 0; i < vec_size / 2; ++i) { + uint32_t y; + uint16_t x = *(uint16_t*)&src[i * 2]; + asm volatile("cvt.rn.f16x2.e4m3x2 %0, %1;" : "=r"(y) : "h"(x)); + *(uint32_t*)&dst[i * 2] = y; + } + } +#else + if constexpr (vec_size == 1) { + dst[0] = half(src[0]); + } else if constexpr (vec_size == 2) { + dst[0] = half(src[0]); + dst[1] = half(src[1]); + } else { + static_assert(vec_size % 4 == 0, "vec_size must be a multiple of 4"); +#pragma unroll + for (uint32_t i = 0; i < vec_size / 4; ++i) { + fast_dequant_f8f16x4<__nv_fp8_e4m3, half>((uint32_t*)&src[i * 4], (uint2*)&dst[i * 4]); + } + } +#endif // FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED + } +}; + +template <> +struct vec_cast { + template + FLASHINFER_INLINE static void cast(half* dst, const __nv_fp8_e5m2* src) { +#ifdef FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED + if constexpr (vec_size == 1) { + dst[0] = half(src[0]); + } else { +#pragma unroll + for (size_t i = 0; i < vec_size / 2; ++i) { + uint32_t y; + uint16_t x = *(uint16_t*)&src[i * 2]; + asm volatile("cvt.rn.f16x2.e5m2x2 %0, %1;" : "=r"(y) : "h"(x)); + *(uint32_t*)&dst[i * 2] = y; + } + } +#else + if constexpr (vec_size == 1) { + dst[0] = half(src[0]); + } else if constexpr (vec_size == 2) { + dst[0] = half(src[0]); + dst[1] = half(src[1]); + } else { + static_assert(vec_size % 4 == 0, "vec_size must be a multiple of 4"); +#pragma unroll + for (uint32_t i = 0; i < vec_size / 4; ++i) { + fast_dequant_f8f16x4<__nv_fp8_e5m2, half>((uint32_t*)&src[i * 4], (uint2*)&dst[i * 4]); + } + } +#endif // FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED + } +}; + +template <> +struct vec_cast { + template + FLASHINFER_INLINE static void cast(float* dst, const nv_bfloat16* src) { + if constexpr (vec_size == 1) { + dst[0] = (float)src[0]; + } else { +#pragma unroll + for (size_t i = 0; i < vec_size / 2; ++i) { + ((float2*)dst)[i] = __bfloat1622float2(((nv_bfloat162*)src)[i]); + } + } + } +}; + +template <> +struct vec_cast { + template + FLASHINFER_INLINE static void cast(nv_bfloat16* dst, const float* src) { + if constexpr (vec_size == 1) { + dst[0] = nv_bfloat16(src[0]); + } else { +#pragma unroll + for (size_t i = 0; i < vec_size / 2; ++i) { + ((nv_bfloat162*)dst)[i] = __float22bfloat162_rn(((float2*)src)[i]); + } + } + } +}; + +template +struct vec_t { + FLASHINFER_INLINE float_t& operator[](size_t i); + FLASHINFER_INLINE const float_t& operator[](size_t i) const; + FLASHINFER_INLINE void fill(float_t val); + FLASHINFER_INLINE void load(const float_t* ptr); + FLASHINFER_INLINE void store(float_t* ptr) const; + FLASHINFER_INLINE void load_global_acquire(float* addr); + FLASHINFER_INLINE void store_global_release(float* addr) const; + FLASHINFER_INLINE void load_global_volatile(float* addr); + FLASHINFER_INLINE void store_global_volatile(float* addr) const; + template + FLASHINFER_INLINE void cast_from(const vec_t& src); + template + FLASHINFER_INLINE void cast_load(const T* ptr); + template + FLASHINFER_INLINE void cast_store(T* ptr) const; + FLASHINFER_INLINE static void memcpy(float_t* dst, const float_t* src); + FLASHINFER_INLINE float_t* ptr(); +}; + +template +FLASHINFER_INLINE void cast_from_impl(vec_t& dst, + const vec_t& src) { + vec_cast::cast( + dst.ptr(), const_cast*>(&src)->ptr()); +} + +template +FLASHINFER_INLINE void cast_load_impl(vec_t& dst, + const src_float_t* src_ptr) { + if constexpr (std::is_same_v) { + dst.load(src_ptr); + } else { + vec_t tmp; + tmp.load(src_ptr); + dst.cast_from(tmp); + } +} + +template +FLASHINFER_INLINE void cast_store_impl(tgt_float_t* dst_ptr, + const vec_t& src) { + if constexpr (std::is_same_v) { + src.store(dst_ptr); + } else { + vec_t tmp; + tmp.cast_from(src); + tmp.store(dst_ptr); + } +} + +/******************* vec_t<__nv_fp8_e4m3> *******************/ + +// __nv_fp8_e4m3 x 1 +template <> +struct vec_t<__nv_fp8_e4m3, 1> { + __nv_fp8_e4m3 data; + + FLASHINFER_INLINE __nv_fp8_e4m3& operator[](size_t i) { return ((__nv_fp8_e4m3*)(&data))[i]; } + FLASHINFER_INLINE const __nv_fp8_e4m3& operator[](size_t i) const { + return ((const __nv_fp8_e4m3*)(&data))[i]; + } + FLASHINFER_INLINE __nv_fp8_e4m3* ptr() { return reinterpret_cast<__nv_fp8_e4m3*>(&data); } + FLASHINFER_INLINE void fill(__nv_fp8_e4m3 val); + FLASHINFER_INLINE void load(const __nv_fp8_e4m3* ptr); + FLASHINFER_INLINE void store(__nv_fp8_e4m3* ptr) const; + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + + FLASHINFER_INLINE static void memcpy(__nv_fp8_e4m3* dst, const __nv_fp8_e4m3* src); +}; + +FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 1>::fill(__nv_fp8_e4m3 val) { data = val; } + +FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 1>::load(const __nv_fp8_e4m3* ptr) { data = *ptr; } + +FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 1>::store(__nv_fp8_e4m3* ptr) const { *ptr = data; } + +FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 1>::memcpy(__nv_fp8_e4m3* dst, + const __nv_fp8_e4m3* src) { + *dst = *src; +} + +// __nv_fp8_e4m3 x 2 +template <> +struct vec_t<__nv_fp8_e4m3, 2> { + __nv_fp8x2_e4m3 data; + + FLASHINFER_INLINE __nv_fp8_e4m3& operator[](size_t i) { return ((__nv_fp8_e4m3*)(&data))[i]; } + FLASHINFER_INLINE const __nv_fp8_e4m3& operator[](size_t i) const { + return ((const __nv_fp8_e4m3*)(&data))[i]; + } + FLASHINFER_INLINE __nv_fp8_e4m3* ptr() { return reinterpret_cast<__nv_fp8_e4m3*>(&data); } + FLASHINFER_INLINE void fill(__nv_fp8_e4m3 val); + FLASHINFER_INLINE void load(const __nv_fp8_e4m3* ptr); + FLASHINFER_INLINE void store(__nv_fp8_e4m3* ptr) const; + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + FLASHINFER_INLINE static void memcpy(__nv_fp8_e4m3* dst, const __nv_fp8_e4m3* src); +}; + +FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 2>::fill(__nv_fp8_e4m3 val) { + data.__x = (__nv_fp8x2_storage_t(val.__x) << 8) | __nv_fp8x2_storage_t(val.__x); +} + +FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 2>::load(const __nv_fp8_e4m3* ptr) { + data = *((__nv_fp8x2_e4m3*)ptr); +} + +FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 2>::store(__nv_fp8_e4m3* ptr) const { + *((__nv_fp8x2_e4m3*)ptr) = data; +} + +FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 2>::memcpy(__nv_fp8_e4m3* dst, + const __nv_fp8_e4m3* src) { + *((__nv_fp8x2_e4m3*)dst) = *((__nv_fp8x2_e4m3*)src); +} + +// __nv_fp8_e4m3 x 4 + +template <> +struct vec_t<__nv_fp8_e4m3, 4> { + __nv_fp8x4_e4m3 data; + + FLASHINFER_INLINE __nv_fp8_e4m3& operator[](size_t i) { return ((__nv_fp8_e4m3*)(&data))[i]; } + FLASHINFER_INLINE const __nv_fp8_e4m3& operator[](size_t i) const { + return ((const __nv_fp8_e4m3*)(&data))[i]; + } + FLASHINFER_INLINE __nv_fp8_e4m3* ptr() { return reinterpret_cast<__nv_fp8_e4m3*>(&data); } + FLASHINFER_INLINE void fill(__nv_fp8_e4m3 val); + FLASHINFER_INLINE void load(const __nv_fp8_e4m3* ptr); + FLASHINFER_INLINE void store(__nv_fp8_e4m3* ptr) const; + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + + FLASHINFER_INLINE static void memcpy(__nv_fp8_e4m3* dst, const __nv_fp8_e4m3* src); +}; + +FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 4>::fill(__nv_fp8_e4m3 val) { + data.__x = (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) | + (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x); +} + +FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 4>::load(const __nv_fp8_e4m3* ptr) { + data = *((__nv_fp8x4_e4m3*)ptr); +} + +FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 4>::store(__nv_fp8_e4m3* ptr) const { + *((__nv_fp8x4_e4m3*)ptr) = data; +} + +FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 4>::memcpy(__nv_fp8_e4m3* dst, + const __nv_fp8_e4m3* src) { + *((__nv_fp8x4_e4m3*)dst) = *((__nv_fp8x4_e4m3*)src); +} + +// __nv_fp8_e4m3 x 8 + +template <> +struct vec_t<__nv_fp8_e4m3, 8> { + uint2 data; + + FLASHINFER_INLINE __nv_fp8_e4m3& operator[](size_t i) { return ((__nv_fp8_e4m3*)(&data))[i]; } + FLASHINFER_INLINE const __nv_fp8_e4m3& operator[](size_t i) const { + return ((const __nv_fp8_e4m3*)(&data))[i]; + } + FLASHINFER_INLINE __nv_fp8_e4m3* ptr() { return reinterpret_cast<__nv_fp8_e4m3*>(&data); } + FLASHINFER_INLINE void fill(__nv_fp8_e4m3 val); + FLASHINFER_INLINE void load(const __nv_fp8_e4m3* ptr); + FLASHINFER_INLINE void store(__nv_fp8_e4m3* ptr) const; + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + + FLASHINFER_INLINE static void memcpy(__nv_fp8_e4m3* dst, const __nv_fp8_e4m3* src); +}; + +FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 8>::fill(__nv_fp8_e4m3 val) { + ((__nv_fp8x4_e4m3*)(&data.x))->__x = + (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) | + (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x); + ((__nv_fp8x4_e4m3*)(&data.y))->__x = + (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) | + (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x); +} + +FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 8>::load(const __nv_fp8_e4m3* ptr) { + data = *((uint2*)ptr); +} + +FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 8>::store(__nv_fp8_e4m3* ptr) const { + *((uint2*)ptr) = data; +} + +FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 8>::memcpy(__nv_fp8_e4m3* dst, + const __nv_fp8_e4m3* src) { + *((uint2*)dst) = *((uint2*)src); +} + +// __nv_fp8_e4m3 x 16 or more +template +struct vec_t<__nv_fp8_e4m3, vec_size> { + static_assert(vec_size % 16 == 0, "Invalid vector size"); + int4 data[vec_size / 16]; + + FLASHINFER_INLINE __nv_fp8_e4m3& operator[](size_t i) { return ((__nv_fp8_e4m3*)data)[i]; } + FLASHINFER_INLINE const __nv_fp8_e4m3& operator[](size_t i) const { + return ((const __nv_fp8_e4m3*)data)[i]; + } + FLASHINFER_INLINE __nv_fp8_e4m3* ptr() { return reinterpret_cast<__nv_fp8_e4m3*>(&data); } + FLASHINFER_INLINE void fill(__nv_fp8_e4m3 val) { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) { + ((__nv_fp8x4_e4m3*)(&(data[i].x)))->__x = + (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) | + (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x); + ((__nv_fp8x4_e4m3*)(&(data[i].y)))->__x = + (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) | + (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x); + ((__nv_fp8x4_e4m3*)(&(data[i].z)))->__x = + (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) | + (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x); + ((__nv_fp8x4_e4m3*)(&(data[i].w)))->__x = + (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) | + (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x); + } + } + FLASHINFER_INLINE void load(const __nv_fp8_e4m3* ptr) { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) { + data[i] = ((int4*)ptr)[i]; + } + } + FLASHINFER_INLINE void store(__nv_fp8_e4m3* ptr) const { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) { + ((int4*)ptr)[i] = data[i]; + } + } + FLASHINFER_INLINE void load_global_acquire(__nv_fp8_e4m3* addr) { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) { + *((int4*)(data + i)) = ld_global_acquire((int4*)(addr + i * 16)); + } + } + FLASHINFER_INLINE void store_global_release(__nv_fp8_e4m3* addr) const { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) { + st_global_release(data[i], (int4*)(addr + i * 16)); + } + } + FLASHINFER_INLINE void load_global_volatile(__nv_fp8_e4m3* addr) { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) { + data[i] = ld_global_volatile((int4*)(addr + i * 16)); + } + } + FLASHINFER_INLINE void store_global_volatile(__nv_fp8_e4m3* addr) const { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) { + st_global_volatile(data[i], (int4*)(addr + i * 16)); + } + } + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + + FLASHINFER_INLINE static void memcpy(__nv_fp8_e4m3* dst, const __nv_fp8_e4m3* src) { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) { + ((int4*)dst)[i] = ((int4*)src)[i]; + } + } +}; + +/******************* vec_t<__nv_fp8_e5m2> *******************/ + +// __nv_fp8_e5m2 x 1 +template <> +struct vec_t<__nv_fp8_e5m2, 1> { + __nv_fp8_e5m2 data; + + FLASHINFER_INLINE __nv_fp8_e5m2& operator[](size_t i) { return ((__nv_fp8_e5m2*)(&data))[i]; } + FLASHINFER_INLINE const __nv_fp8_e5m2& operator[](size_t i) const { + return ((const __nv_fp8_e5m2*)(&data))[i]; + } + FLASHINFER_INLINE __nv_fp8_e5m2* ptr() { return reinterpret_cast<__nv_fp8_e5m2*>(&data); } + FLASHINFER_INLINE void fill(__nv_fp8_e5m2 val); + FLASHINFER_INLINE void load(const __nv_fp8_e5m2* ptr); + FLASHINFER_INLINE void store(__nv_fp8_e5m2* ptr) const; + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + + FLASHINFER_INLINE static void memcpy(__nv_fp8_e5m2* dst, const __nv_fp8_e5m2* src); +}; + +FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 1>::fill(__nv_fp8_e5m2 val) { data = val; } + +FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 1>::load(const __nv_fp8_e5m2* ptr) { data = *ptr; } + +FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 1>::store(__nv_fp8_e5m2* ptr) const { *ptr = data; } + +FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 1>::memcpy(__nv_fp8_e5m2* dst, + const __nv_fp8_e5m2* src) { + *dst = *src; +} + +// __nv_fp8_e5m2 x 2 +template <> +struct vec_t<__nv_fp8_e5m2, 2> { + __nv_fp8x2_e5m2 data; + + FLASHINFER_INLINE __nv_fp8_e5m2& operator[](size_t i) { return ((__nv_fp8_e5m2*)(&data))[i]; } + FLASHINFER_INLINE const __nv_fp8_e5m2& operator[](size_t i) const { + return ((const __nv_fp8_e5m2*)(&data))[i]; + } + FLASHINFER_INLINE __nv_fp8_e5m2* ptr() { return reinterpret_cast<__nv_fp8_e5m2*>(&data); } + FLASHINFER_INLINE void fill(__nv_fp8_e5m2 val); + FLASHINFER_INLINE void load(const __nv_fp8_e5m2* ptr); + FLASHINFER_INLINE void store(__nv_fp8_e5m2* ptr) const; + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + + FLASHINFER_INLINE static void memcpy(__nv_fp8_e5m2* dst, const __nv_fp8_e5m2* src); +}; + +FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 2>::fill(__nv_fp8_e5m2 val) { + data.__x = (__nv_fp8x2_storage_t(val.__x) << 8) | __nv_fp8x2_storage_t(val.__x); +} + +FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 2>::load(const __nv_fp8_e5m2* ptr) { + data = *((__nv_fp8x2_e5m2*)ptr); +} + +FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 2>::store(__nv_fp8_e5m2* ptr) const { + *((__nv_fp8x2_e5m2*)ptr) = data; +} + +FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 2>::memcpy(__nv_fp8_e5m2* dst, + const __nv_fp8_e5m2* src) { + *((__nv_fp8x2_e5m2*)dst) = *((__nv_fp8x2_e5m2*)src); +} + +// __nv_fp8_e5m2 x 4 + +template <> +struct vec_t<__nv_fp8_e5m2, 4> { + __nv_fp8x4_e5m2 data; + + FLASHINFER_INLINE __nv_fp8_e5m2& operator[](size_t i) { return ((__nv_fp8_e5m2*)(&data))[i]; } + FLASHINFER_INLINE const __nv_fp8_e5m2& operator[](size_t i) const { + return ((const __nv_fp8_e5m2*)(&data))[i]; + } + FLASHINFER_INLINE __nv_fp8_e5m2* ptr() { return reinterpret_cast<__nv_fp8_e5m2*>(&data); } + FLASHINFER_INLINE void fill(__nv_fp8_e5m2 val); + FLASHINFER_INLINE void load(const __nv_fp8_e5m2* ptr); + FLASHINFER_INLINE void store(__nv_fp8_e5m2* ptr) const; + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + + FLASHINFER_INLINE static void memcpy(__nv_fp8_e5m2* dst, const __nv_fp8_e5m2* src); +}; + +FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 4>::fill(__nv_fp8_e5m2 val) { + data.__x = (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) | + (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x); +} + +FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 4>::load(const __nv_fp8_e5m2* ptr) { + data = *((__nv_fp8x4_e5m2*)ptr); +} + +FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 4>::store(__nv_fp8_e5m2* ptr) const { + *((__nv_fp8x4_e5m2*)ptr) = data; +} + +FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 4>::memcpy(__nv_fp8_e5m2* dst, + const __nv_fp8_e5m2* src) { + *((__nv_fp8x4_e5m2*)dst) = *((__nv_fp8x4_e5m2*)src); +} + +// __nv_fp8_e5m2 x 8 + +template <> +struct vec_t<__nv_fp8_e5m2, 8> { + uint2 data; + + FLASHINFER_INLINE __nv_fp8_e5m2& operator[](size_t i) { return ((__nv_fp8_e5m2*)(&data))[i]; } + FLASHINFER_INLINE const __nv_fp8_e5m2& operator[](size_t i) const { + return ((const __nv_fp8_e5m2*)(&data))[i]; + } + FLASHINFER_INLINE __nv_fp8_e5m2* ptr() { return reinterpret_cast<__nv_fp8_e5m2*>(&data); } + FLASHINFER_INLINE void fill(__nv_fp8_e5m2 val); + FLASHINFER_INLINE void load(const __nv_fp8_e5m2* ptr); + FLASHINFER_INLINE void store(__nv_fp8_e5m2* ptr) const; + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + FLASHINFER_INLINE static void memcpy(__nv_fp8_e5m2* dst, const __nv_fp8_e5m2* src); +}; + +FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 8>::fill(__nv_fp8_e5m2 val) { + ((__nv_fp8x4_e5m2*)(&data.x))->__x = + (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) | + (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x); + ((__nv_fp8x4_e5m2*)(&data.y))->__x = + (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) | + (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x); +} + +FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 8>::load(const __nv_fp8_e5m2* ptr) { + data = *((uint2*)ptr); +} + +FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 8>::store(__nv_fp8_e5m2* ptr) const { + *((uint2*)ptr) = data; +} + +FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 8>::memcpy(__nv_fp8_e5m2* dst, + const __nv_fp8_e5m2* src) { + *((uint2*)dst) = *((uint2*)src); +} + +// __nv_fp8_e5m2 x 16 or more + +template +struct vec_t<__nv_fp8_e5m2, vec_size> { + static_assert(vec_size % 16 == 0, "Invalid vector size"); + int4 data[vec_size / 16]; + + FLASHINFER_INLINE __nv_fp8_e5m2& operator[](size_t i) { return ((__nv_fp8_e5m2*)data)[i]; } + FLASHINFER_INLINE const __nv_fp8_e5m2& operator[](size_t i) const { + return ((const __nv_fp8_e5m2*)data)[i]; + } + FLASHINFER_INLINE __nv_fp8_e5m2* ptr() { return reinterpret_cast<__nv_fp8_e5m2*>(&data); } + FLASHINFER_INLINE void fill(__nv_fp8_e5m2 val) { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) { + ((__nv_fp8x4_e5m2*)(&(data[i].x)))->__x = + (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) | + (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x); + ((__nv_fp8x4_e5m2*)(&(data[i].y)))->__x = + (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) | + (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x); + ((__nv_fp8x4_e5m2*)(&(data[i].z)))->__x = + (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) | + (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x); + ((__nv_fp8x4_e5m2*)(&(data[i].w)))->__x = + (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) | + (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x); + } + } + FLASHINFER_INLINE void load(const __nv_fp8_e5m2* ptr) { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) { + data[i] = ((int4*)ptr)[i]; + } + } + FLASHINFER_INLINE void store(__nv_fp8_e5m2* ptr) const { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) { + ((int4*)ptr)[i] = data[i]; + } + } + FLASHINFER_INLINE void store_global_release(__nv_fp8_e5m2* addr) const { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) { + st_global_release(data[i], (int4*)(addr + i * 16)); + } + } + FLASHINFER_INLINE void load_global_acquire(__nv_fp8_e5m2* addr) { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) { + data[i] = ld_global_acquire((int4*)(addr + i * 16)); + } + } + FLASHINFER_INLINE void store_global_volatile(__nv_fp8_e5m2* addr) const { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) { + st_global_volatile(data[i], (int4*)(addr + i * 16)); + } + } + FLASHINFER_INLINE void load_global_volatile(__nv_fp8_e5m2* addr) { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) { + data[i] = ld_global_volatile((int4*)(addr + i * 16)); + } + } + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + FLASHINFER_INLINE static void memcpy(__nv_fp8_e5m2* dst, const __nv_fp8_e5m2* src) { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) { + ((int4*)dst)[i] = ((int4*)src)[i]; + } + } +}; + +/******************* vec_t *******************/ + +// half x 1 +template <> +struct vec_t { + half data; + + FLASHINFER_INLINE half& operator[](size_t i) { return ((half*)(&data))[i]; } + FLASHINFER_INLINE const half& operator[](size_t i) const { return ((const half*)(&data))[i]; } + FLASHINFER_INLINE half* ptr() { return reinterpret_cast(&data); } + FLASHINFER_INLINE void fill(half val); + FLASHINFER_INLINE void load(const half* ptr); + FLASHINFER_INLINE void store(half* ptr) const; + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + + FLASHINFER_INLINE static void memcpy(half* dst, const half* src); +}; + +FLASHINFER_INLINE void vec_t::fill(half val) { data = val; } + +FLASHINFER_INLINE void vec_t::load(const half* ptr) { data = *ptr; } + +FLASHINFER_INLINE void vec_t::store(half* ptr) const { *ptr = data; } + +FLASHINFER_INLINE void vec_t::memcpy(half* dst, const half* src) { *dst = *src; } + +// half x 2 +template <> +struct vec_t { + half2 data; + + FLASHINFER_INLINE half& operator[](size_t i) { return ((half*)(&data))[i]; } + FLASHINFER_INLINE const half& operator[](size_t i) const { return ((const half*)(&data))[i]; } + FLASHINFER_INLINE half* ptr() { return reinterpret_cast(&data); } + FLASHINFER_INLINE void fill(half val); + FLASHINFER_INLINE void load(const half* ptr); + FLASHINFER_INLINE void store(half* ptr) const; + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + + FLASHINFER_INLINE static void memcpy(half* dst, const half* src); +}; + +FLASHINFER_INLINE void vec_t::fill(half val) { data = make_half2(val, val); } + +FLASHINFER_INLINE void vec_t::load(const half* ptr) { data = *((half2*)ptr); } + +FLASHINFER_INLINE void vec_t::store(half* ptr) const { *((half2*)ptr) = data; } + +FLASHINFER_INLINE void vec_t::memcpy(half* dst, const half* src) { + *((half2*)dst) = *((half2*)src); +} + +// half x 4 + +template <> +struct vec_t { + uint2 data; + + FLASHINFER_INLINE half& operator[](size_t i) { return ((half*)(&data))[i]; } + FLASHINFER_INLINE const half& operator[](size_t i) const { return ((const half*)(&data))[i]; } + FLASHINFER_INLINE half* ptr() { return reinterpret_cast(&data); } + FLASHINFER_INLINE void fill(half val); + FLASHINFER_INLINE void load(const half* ptr); + FLASHINFER_INLINE void store(half* ptr) const; + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + FLASHINFER_INLINE static void memcpy(half* dst, const half* src); +}; + +FLASHINFER_INLINE void vec_t::fill(half val) { + *(half2*)(&data.x) = make_half2(val, val); + *(half2*)(&data.y) = make_half2(val, val); +} + +FLASHINFER_INLINE void vec_t::load(const half* ptr) { data = *((uint2*)ptr); } + +FLASHINFER_INLINE void vec_t::store(half* ptr) const { *((uint2*)ptr) = data; } + +FLASHINFER_INLINE void vec_t::memcpy(half* dst, const half* src) { + *((uint2*)dst) = *((uint2*)src); +} + +// half x 8 or more + +template +struct vec_t { + static_assert(vec_size % 8 == 0, "Invalid vector size"); + int4 data[vec_size / 8]; + FLASHINFER_INLINE half& operator[](size_t i) { return ((half*)data)[i]; } + FLASHINFER_INLINE const half& operator[](size_t i) const { return ((const half*)data)[i]; } + FLASHINFER_INLINE half* ptr() { return reinterpret_cast(&data); } + FLASHINFER_INLINE void fill(half val) { +#pragma unroll + for (size_t i = 0; i < vec_size / 8; ++i) { + *(half2*)(&(data[i].x)) = make_half2(val, val); + *(half2*)(&(data[i].y)) = make_half2(val, val); + *(half2*)(&(data[i].z)) = make_half2(val, val); + *(half2*)(&(data[i].w)) = make_half2(val, val); + } + } + FLASHINFER_INLINE void load(const half* ptr) { +#pragma unroll + for (size_t i = 0; i < vec_size / 8; ++i) { + data[i] = ((int4*)ptr)[i]; + } + } + FLASHINFER_INLINE void store(half* ptr) const { +#pragma unroll + for (size_t i = 0; i < vec_size / 8; ++i) { + ((int4*)ptr)[i] = data[i]; + } + } + FLASHINFER_INLINE void load_global_acquire(half* addr) { +#pragma unroll + for (size_t i = 0; i < vec_size / 8; ++i) { + data[i] = ld_global_acquire((int4*)(addr + i * 8)); + } + } + FLASHINFER_INLINE void store_global_release(half* addr) const { +#pragma unroll + for (size_t i = 0; i < vec_size / 8; ++i) { + st_global_release(data[i], (int4*)(addr + i * 8)); + } + } + FLASHINFER_INLINE void store_global_volatile(half* addr) const { +#pragma unroll + for (size_t i = 0; i < vec_size / 8; ++i) { + st_global_volatile(data[i], (int4*)(addr + i * 8)); + } + } + FLASHINFER_INLINE void load_global_volatile(half* addr) { +#pragma unroll + for (size_t i = 0; i < vec_size / 8; ++i) { + data[i] = ld_global_volatile((int4*)(addr + i * 8)); + } + } + + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + FLASHINFER_INLINE static void memcpy(half* dst, const half* src) { +#pragma unroll + for (size_t i = 0; i < vec_size / 8; ++i) { + ((int4*)dst)[i] = ((int4*)src)[i]; + } + } +}; + +/******************* vec_t *******************/ + +// nv_bfloat16 x 1 +template <> +struct vec_t { + nv_bfloat16 data; + FLASHINFER_INLINE nv_bfloat16& operator[](size_t i) { return ((nv_bfloat16*)(&data))[i]; } + FLASHINFER_INLINE const nv_bfloat16& operator[](size_t i) const { + return ((const nv_bfloat16*)(&data))[i]; + } + FLASHINFER_INLINE nv_bfloat16* ptr() { return reinterpret_cast(&data); } + FLASHINFER_INLINE void fill(nv_bfloat16 val); + FLASHINFER_INLINE void load(const nv_bfloat16* ptr); + FLASHINFER_INLINE void store(nv_bfloat16* ptr) const; + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + FLASHINFER_INLINE static void memcpy(nv_bfloat16* dst, const nv_bfloat16* src); +}; + +FLASHINFER_INLINE void vec_t::fill(nv_bfloat16 val) { data = val; } + +FLASHINFER_INLINE void vec_t::load(const nv_bfloat16* ptr) { data = *ptr; } + +FLASHINFER_INLINE void vec_t::store(nv_bfloat16* ptr) const { *ptr = data; } + +FLASHINFER_INLINE void vec_t::memcpy(nv_bfloat16* dst, const nv_bfloat16* src) { + *dst = *src; +} + +// nv_bfloat16 x 2 +template <> +struct vec_t { + nv_bfloat162 data; + + FLASHINFER_INLINE nv_bfloat16& operator[](size_t i) { return ((nv_bfloat16*)(&data))[i]; } + FLASHINFER_INLINE const nv_bfloat16& operator[](size_t i) const { + return ((const nv_bfloat16*)(&data))[i]; + } + FLASHINFER_INLINE nv_bfloat16* ptr() { return reinterpret_cast(&data); } + FLASHINFER_INLINE void fill(nv_bfloat16 val); + FLASHINFER_INLINE void load(const nv_bfloat16* ptr); + FLASHINFER_INLINE void store(nv_bfloat16* ptr) const; + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + FLASHINFER_INLINE static void memcpy(nv_bfloat16* dst, const nv_bfloat16* src); +}; + +FLASHINFER_INLINE void vec_t::fill(nv_bfloat16 val) { + data = make_bfloat162(val, val); +} + +FLASHINFER_INLINE void vec_t::load(const nv_bfloat16* ptr) { + data = *((nv_bfloat162*)ptr); +} + +FLASHINFER_INLINE void vec_t::store(nv_bfloat16* ptr) const { + *((nv_bfloat162*)ptr) = data; +} + +FLASHINFER_INLINE void vec_t::memcpy(nv_bfloat16* dst, const nv_bfloat16* src) { + *((nv_bfloat162*)dst) = *((nv_bfloat162*)src); +} + +// nv_bfloat16 x 4 + +template <> +struct vec_t { + uint2 data; + + FLASHINFER_INLINE nv_bfloat16& operator[](size_t i) { return ((nv_bfloat16*)(&data))[i]; } + FLASHINFER_INLINE const nv_bfloat16& operator[](size_t i) const { + return ((const nv_bfloat16*)(&data))[i]; + } + FLASHINFER_INLINE nv_bfloat16* ptr() { return reinterpret_cast(&data); } + FLASHINFER_INLINE void fill(nv_bfloat16 val); + FLASHINFER_INLINE void load(const nv_bfloat16* ptr); + FLASHINFER_INLINE void store(nv_bfloat16* ptr) const; + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + FLASHINFER_INLINE static void memcpy(nv_bfloat16* dst, const nv_bfloat16* src); +}; + +FLASHINFER_INLINE void vec_t::fill(nv_bfloat16 val) { + *(nv_bfloat162*)(&data.x) = make_bfloat162(val, val); + *(nv_bfloat162*)(&data.y) = make_bfloat162(val, val); +} + +FLASHINFER_INLINE void vec_t::load(const nv_bfloat16* ptr) { + data = *((uint2*)ptr); +} + +FLASHINFER_INLINE void vec_t::store(nv_bfloat16* ptr) const { + *((uint2*)ptr) = data; +} + +FLASHINFER_INLINE void vec_t::memcpy(nv_bfloat16* dst, const nv_bfloat16* src) { + *((uint2*)dst) = *((uint2*)src); +} + +// nv_bfloat16 x 8 or more + +template +struct vec_t { + static_assert(vec_size % 8 == 0, "Invalid vector size"); + int4 data[vec_size / 8]; + + FLASHINFER_INLINE nv_bfloat16& operator[](size_t i) { return ((nv_bfloat16*)data)[i]; } + FLASHINFER_INLINE const nv_bfloat16& operator[](size_t i) const { + return ((const nv_bfloat16*)data)[i]; + } + FLASHINFER_INLINE nv_bfloat16* ptr() { return reinterpret_cast(&data); } + FLASHINFER_INLINE void fill(nv_bfloat16 val) { +#pragma unoll + for (size_t i = 0; i < vec_size / 8; ++i) { + *(nv_bfloat162*)(&(data[i].x)) = make_bfloat162(val, val); + *(nv_bfloat162*)(&(data[i].y)) = make_bfloat162(val, val); + *(nv_bfloat162*)(&(data[i].z)) = make_bfloat162(val, val); + *(nv_bfloat162*)(&(data[i].w)) = make_bfloat162(val, val); + } + } + FLASHINFER_INLINE void load(const nv_bfloat16* ptr) { +#pragma unoll + for (size_t i = 0; i < vec_size / 8; ++i) { + data[i] = ((int4*)ptr)[i]; + } + } + FLASHINFER_INLINE void store(nv_bfloat16* ptr) const { +#pragma unoll + for (size_t i = 0; i < vec_size / 8; ++i) { + ((int4*)ptr)[i] = data[i]; + } + } + FLASHINFER_INLINE void store_global_release(nv_bfloat16* addr) const { +#pragma unoll + for (size_t i = 0; i < vec_size / 8; ++i) { + st_global_release(data[i], (int4*)(addr + i * 8)); + } + } + FLASHINFER_INLINE void load_global_acquire(nv_bfloat16* addr) { +#pragma unoll + for (size_t i = 0; i < vec_size / 8; ++i) { + data[i] = ld_global_acquire((int4*)(addr + i * 8)); + } + } + FLASHINFER_INLINE void store_global_volatile(nv_bfloat16* addr) const { +#pragma unoll + for (size_t i = 0; i < vec_size / 8; ++i) { + st_global_volatile(data[i], (int4*)(addr + i * 8)); + } + } + FLASHINFER_INLINE void load_global_volatile(nv_bfloat16* addr) { +#pragma unoll + for (size_t i = 0; i < vec_size / 8; ++i) { + data[i] = ld_global_volatile((int4*)(addr + i * 8)); + } + } + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + FLASHINFER_INLINE static void memcpy(nv_bfloat16* dst, const nv_bfloat16* src) { +#pragma unoll + for (size_t i = 0; i < vec_size / 8; ++i) { + ((int4*)dst)[i] = ((int4*)src)[i]; + } + } +}; + +/******************* vec_t *******************/ + +// uint8_t x 1 +template <> +struct vec_t { + uint8_t data; + + FLASHINFER_INLINE uint8_t& operator[](size_t i) { return ((uint8_t*)(&data))[i]; } + FLASHINFER_INLINE const uint8_t& operator[](size_t i) const { + return ((const uint8_t*)(&data))[i]; + } + FLASHINFER_INLINE uint8_t* ptr() { return reinterpret_cast(&data); } + FLASHINFER_INLINE void fill(uint8_t val); + FLASHINFER_INLINE void load(const uint8_t* ptr); + FLASHINFER_INLINE void store(uint8_t* ptr) const; + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + + FLASHINFER_INLINE static void memcpy(uint8_t* dst, const uint8_t* src); +}; + +FLASHINFER_INLINE void vec_t::fill(uint8_t val) { data = val; } + +FLASHINFER_INLINE void vec_t::load(const uint8_t* ptr) { data = *ptr; } + +FLASHINFER_INLINE void vec_t::store(uint8_t* ptr) const { *ptr = data; } + +FLASHINFER_INLINE void vec_t::memcpy(uint8_t* dst, const uint8_t* src) { *dst = *src; } + +// uint8_t x 2 +template <> +struct vec_t { + uint16_t data; + + FLASHINFER_INLINE uint8_t& operator[](size_t i) { return ((uint8_t*)(&data))[i]; } + FLASHINFER_INLINE const uint8_t& operator[](size_t i) const { + return ((const uint8_t*)(&data))[i]; + } + FLASHINFER_INLINE uint8_t* ptr() { return reinterpret_cast(&data); } + FLASHINFER_INLINE void fill(uint8_t val); + FLASHINFER_INLINE void load(const uint8_t* ptr); + FLASHINFER_INLINE void store(uint8_t* ptr) const; + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + + FLASHINFER_INLINE static void memcpy(uint8_t* dst, const uint8_t* src); +}; + +FLASHINFER_INLINE void vec_t::fill(uint8_t val) { + data = (uint16_t(val) << 8) | uint16_t(val); +} + +FLASHINFER_INLINE void vec_t::load(const uint8_t* ptr) { data = *((uint16_t*)ptr); } + +FLASHINFER_INLINE void vec_t::store(uint8_t* ptr) const { *((uint16_t*)ptr) = data; } + +FLASHINFER_INLINE void vec_t::memcpy(uint8_t* dst, const uint8_t* src) { + *((uint16_t*)dst) = *((uint16_t*)src); +} + +// uint8_t x 4 + +template <> +struct vec_t { + uint32_t data; + + FLASHINFER_INLINE uint8_t& operator[](size_t i) { return ((uint8_t*)(&data))[i]; } + FLASHINFER_INLINE const uint8_t& operator[](size_t i) const { + return ((const uint8_t*)(&data))[i]; + } + FLASHINFER_INLINE uint8_t* ptr() { return reinterpret_cast(&data); } + FLASHINFER_INLINE void fill(uint8_t val); + FLASHINFER_INLINE void load(const uint8_t* ptr); + FLASHINFER_INLINE void store(uint8_t* ptr) const; + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + + FLASHINFER_INLINE static void memcpy(uint8_t* dst, const uint8_t* src); +}; + +FLASHINFER_INLINE void vec_t::fill(uint8_t val) { + data = (uint32_t(val) << 24) | (uint32_t(val) << 16) | (uint32_t(val) << 8) | uint32_t(val); +} + +FLASHINFER_INLINE void vec_t::load(const uint8_t* ptr) { data = *((uint32_t*)ptr); } + +FLASHINFER_INLINE void vec_t::store(uint8_t* ptr) const { *((uint32_t*)ptr) = data; } + +FLASHINFER_INLINE void vec_t::memcpy(uint8_t* dst, const uint8_t* src) { + *((uint32_t*)dst) = *((uint32_t*)src); +} + +// uint8_t x 8 + +template <> +struct vec_t { + uint2 data; + + FLASHINFER_INLINE uint8_t& operator[](size_t i) { return ((uint8_t*)(&data))[i]; } + FLASHINFER_INLINE const uint8_t& operator[](size_t i) const { + return ((const uint8_t*)(&data))[i]; + } + FLASHINFER_INLINE uint8_t* ptr() { return reinterpret_cast(&data); } + FLASHINFER_INLINE void fill(uint8_t val); + FLASHINFER_INLINE void load(const uint8_t* ptr); + FLASHINFER_INLINE void store(uint8_t* ptr) const; + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + FLASHINFER_INLINE static void memcpy(uint8_t* dst, const uint8_t* src); +}; + +FLASHINFER_INLINE void vec_t::fill(uint8_t val) { + uint32_t val32 = + (uint32_t(val) << 24) | (uint32_t(val) << 16) | (uint32_t(val) << 8) | uint32_t(val); + data.x = val32; + data.y = val32; +} + +FLASHINFER_INLINE void vec_t::load(const uint8_t* ptr) { data = *((uint2*)ptr); } + +FLASHINFER_INLINE void vec_t::store(uint8_t* ptr) const { *((uint2*)ptr) = data; } + +FLASHINFER_INLINE void vec_t::memcpy(uint8_t* dst, const uint8_t* src) { + *((uint2*)dst) = *((uint2*)src); +} + +// uint8_t x 16 or more + +template +struct vec_t { + static_assert(vec_size % 16 == 0, "Invalid vector size"); + int4 data[vec_size / 16]; + + FLASHINFER_INLINE uint8_t& operator[](size_t i) { return ((uint8_t*)data)[i]; } + FLASHINFER_INLINE const uint8_t& operator[](size_t i) const { return ((const uint8_t*)data)[i]; } + FLASHINFER_INLINE uint8_t* ptr() { return reinterpret_cast(&data); } + FLASHINFER_INLINE void fill(uint8_t val) { + uint32_t val32 = + (uint32_t(val) << 24) | (uint32_t(val) << 16) | (uint32_t(val) << 8) | uint32_t(val); +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) { + data[i].x = val32; + data[i].y = val32; + data[i].z = val32; + data[i].w = val32; + } + } + FLASHINFER_INLINE void load(const uint8_t* ptr) { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) { + data[i] = ((int4*)ptr)[i]; + } + } + FLASHINFER_INLINE void store(uint8_t* ptr) const { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) { + ((int4*)ptr)[i] = data[i]; + } + } + FLASHINFER_INLINE void load_global_acquire(uint8_t* addr) { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) { + data[i] = ld_global_acquire((int4*)(addr + i * 16)); + } + } + FLASHINFER_INLINE void store_global_release(uint8_t* addr) const { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) { + st_global_release(data[i], (int4*)(addr + i * 16)); + } + } + FLASHINFER_INLINE void load_global_volatile(uint8_t* addr) { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) { + data[i] = ld_global_volatile((int4*)(addr + i * 16)); + } + } + FLASHINFER_INLINE void store_global_volatile(uint8_t* addr) const { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) { + st_global_volatile(data[i], (int4*)(addr + i * 16)); + } + } + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + FLASHINFER_INLINE static void memcpy(uint8_t* dst, const uint8_t* src) { +#pragma unroll + for (size_t i = 0; i < vec_size / 16; ++i) { + ((int4*)dst)[i] = ((int4*)src)[i]; + } + } +}; + +/******************* vec_t *******************/ + +// float x 1 + +template <> +struct vec_t { + float data; + + FLASHINFER_INLINE float& operator[](size_t i) { return ((float*)(&data))[i]; } + FLASHINFER_INLINE const float& operator[](size_t i) const { return ((const float*)(&data))[i]; } + FLASHINFER_INLINE float* ptr() { return reinterpret_cast(&data); } + FLASHINFER_INLINE void fill(float val); + FLASHINFER_INLINE void load(const float* ptr); + FLASHINFER_INLINE void store(float* ptr) const; + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + FLASHINFER_INLINE static void memcpy(float* dst, const float* src); +}; + +FLASHINFER_INLINE void vec_t::fill(float val) { data = val; } + +FLASHINFER_INLINE void vec_t::load(const float* ptr) { data = *ptr; } + +FLASHINFER_INLINE void vec_t::store(float* ptr) const { *ptr = data; } + +FLASHINFER_INLINE void vec_t::memcpy(float* dst, const float* src) { *dst = *src; } + +// float x 2 + +template <> +struct vec_t { + float2 data; + + FLASHINFER_INLINE float& operator[](size_t i) { return ((float*)(&data))[i]; } + FLASHINFER_INLINE const float& operator[](size_t i) const { return ((const float*)(&data))[i]; } + FLASHINFER_INLINE float* ptr() { return reinterpret_cast(&data); } + FLASHINFER_INLINE void fill(float val); + FLASHINFER_INLINE void load(const float* ptr); + FLASHINFER_INLINE void store(float* ptr) const; + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + FLASHINFER_INLINE static void memcpy(float* dst, const float* src); +}; + +FLASHINFER_INLINE void vec_t::fill(float val) { data = make_float2(val, val); } + +FLASHINFER_INLINE void vec_t::load(const float* ptr) { data = *((float2*)ptr); } + +FLASHINFER_INLINE void vec_t::store(float* ptr) const { *((float2*)ptr) = data; } + +FLASHINFER_INLINE void vec_t::memcpy(float* dst, const float* src) { + *((float2*)dst) = *((float2*)src); +} + +// float x 4 or more +template +struct vec_t { + static_assert(vec_size % 4 == 0, "Invalid vector size"); + float4 data[vec_size / 4]; + + FLASHINFER_INLINE float& operator[](size_t i) { return ((float*)(data))[i]; } + FLASHINFER_INLINE const float& operator[](size_t i) const { return ((const float*)(data))[i]; } + FLASHINFER_INLINE float* ptr() { return reinterpret_cast(&data); } + FLASHINFER_INLINE void fill(float val) { +#pragma unroll + for (size_t i = 0; i < vec_size / 4; ++i) { + data[i] = make_float4(val, val, val, val); + } + } + FLASHINFER_INLINE void load(const float* ptr) { +#pragma unroll + for (size_t i = 0; i < vec_size / 4; ++i) { + data[i] = ((float4*)ptr)[i]; + } + } + FLASHINFER_INLINE void store(float* ptr) const { +#pragma unroll + for (size_t i = 0; i < vec_size / 4; ++i) { + ((float4*)ptr)[i] = data[i]; + } + } + FLASHINFER_INLINE void store_global_release(float* addr) const { +#pragma unroll + for (size_t i = 0; i < vec_size / 4; ++i) { + st_global_release(*(int4*)(data + i), (int4*)(addr + i * 4)); + } + } + FLASHINFER_INLINE void load_global_acquire(float* addr) { +#pragma unroll + for (size_t i = 0; i < vec_size / 4; ++i) { + *((int4*)(data + i)) = ld_global_acquire((int4*)(addr + i * 4)); + } + } + FLASHINFER_INLINE void store_global_volatile(float* addr) const { +#pragma unroll + for (size_t i = 0; i < vec_size / 4; ++i) { + st_global_volatile(*(int4*)(data + i), (int4*)(addr + i * 4)); + } + } + FLASHINFER_INLINE void load_global_volatile(float* addr) { +#pragma unroll + for (size_t i = 0; i < vec_size / 4; ++i) { + *((int4*)(data + i)) = ld_global_volatile((int4*)(addr + i * 4)); + } + } + template + FLASHINFER_INLINE void cast_from(const vec_t& src) { + cast_from_impl(*this, src); + } + template + FLASHINFER_INLINE void cast_load(const T* ptr) { + cast_load_impl(*this, ptr); + } + template + FLASHINFER_INLINE void cast_store(T* ptr) const { + cast_store_impl(ptr, *this); + } + FLASHINFER_INLINE static void memcpy(float* dst, const float* src) { +#pragma unroll + for (size_t i = 0; i < vec_size / 4; ++i) { + ((float4*)dst)[i] = ((float4*)src)[i]; + } + } +}; + +template +struct vec2_dtype { + using type = T; +}; + +template <> +struct vec2_dtype { + using type = half2; +}; + +template <> +struct vec2_dtype<__nv_bfloat16> { + using type = __nv_bfloat162; +}; + +template <> +struct vec2_dtype<__nv_fp8_e4m3> { + using type = __nv_fp8x2_e4m3; +}; + +template <> +struct vec2_dtype<__nv_fp8_e5m2> { + using type = __nv_fp8x2_e5m2; +}; + +template +using vec2_dtype_t = typename vec2_dtype::type; + +template +FLASHINFER_INLINE vec2_dtype_t get_vec2_element(vec_t& vec, int i) { + static_assert(VEC_SIZE % 2 == 0, "VEC_SIZE must be a multiple of 2"); + return ((vec2_dtype_t*)&(vec[0]))[i]; +} + +} // namespace flashinfer + +#endif // VEC_DTYPES_CUH_ diff --git a/include/pytorch_conversion_utils.h b/include/pytorch_conversion_utils.h new file mode 100644 index 0000000000000000000000000000000000000000..abfe18207a1fa84e1edc347e9c8ac49193bb2223 --- /dev/null +++ b/include/pytorch_conversion_utils.h @@ -0,0 +1,29 @@ +/* + * Copyright (c) 2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once +#include + +inline at::Tensor vec_to_tensor(const std::vector& vec) { + return at::tensor(vec, at::dtype(at::kLong).device(at::kCPU)); +} + +inline std::vector tensor_to_vec(const at::Tensor& tensor) { + const size_t size = tensor.numel(); + const int64_t* first = tensor.const_data_ptr(); + const int64_t* last = first + size; + return std::vector(first, last); +} diff --git a/include/pytorch_extension_utils.h b/include/pytorch_extension_utils.h new file mode 100644 index 0000000000000000000000000000000000000000..e2757901b962480fea87027685b26a45c949220c --- /dev/null +++ b/include/pytorch_extension_utils.h @@ -0,0 +1,336 @@ +/* + * Copyright (c) 2023 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once +#include +#include +#include +#include + +#ifdef FLASHINFER_ENABLE_BF16 +#include +#endif + +#ifdef FLASHINFER_ENABLE_F16 +#include +#endif + +#if defined(FLASHINFER_ENABLE_FP8_E4M3) || defined(FLASHINFER_ENABLE_FP8_E5M2) || \ + defined(FLASHINFER_ENABLE_FP8_E8M0) +#include +#endif + +#if defined(FLASHINFER_ENABLE_FP4_E2M1) +#include +#endif + +#ifndef FLASHINFER_EXT_MODULE_INITED +#define FLASHINFER_EXT_MODULE_INITED + +// To expand macros in #name +#define FLASHINFER_EXT_MODULE_INIT_EXPAND(name) FLASHINFER_EXT_MODULE_INIT(name) + +/* Creates a dummy empty module that can be imported from Python. + The import from Python will load the .so consisting of the file + in this extension, so that the TORCH_LIBRARY_FRAGMENT static initializers + are run. */ +#define FLASHINFER_EXT_MODULE_INIT(name) \ + extern "C" { \ + __attribute__((weak)) PyObject* PyInit_##name(void) { \ + static struct PyModuleDef module_def = { \ + PyModuleDef_HEAD_INIT, \ + #name, /* name of module */ \ + NULL, /* module documentation, may be NULL */ \ + -1, /* size of per-interpreter state of the module, \ + or -1 if the module keeps state in global variables. */ \ + NULL, /* methods */ \ + NULL, /* slots */ \ + NULL, /* traverse */ \ + NULL, /* clear */ \ + NULL, /* free */ \ + }; \ + return PyModule_Create(&module_def); \ + } \ + } + +FLASHINFER_EXT_MODULE_INIT_EXPAND(TORCH_EXTENSION_NAME) + +#undef FLASHINFER_EXT_MODULE_INIT +#undef FLASHINFER_EXT_MODULE_INIT_EXPAND + +#endif + +#define _DISPATCH_CASE_I32(c_type, ...) \ + case at::ScalarType::Int: { \ + using c_type = int32_t; \ + return __VA_ARGS__(); \ + } + +#define _DISPATCH_CASE_I64(c_type, ...) \ + case at::ScalarType::Long: { \ + using c_type = int64_t; \ + return __VA_ARGS__(); \ + } + +#define DISPATCH_PYTORCH_IDTYPE_TO_CTYPE(pytorch_dtype, c_type, ...) \ + [&]() -> bool { \ + switch (pytorch_dtype) { \ + _DISPATCH_CASE_I32(c_type, __VA_ARGS__) \ + _DISPATCH_CASE_I64(c_type, __VA_ARGS__) \ + default: \ + std::ostringstream oss; \ + oss << __PRETTY_FUNCTION__ << " failed to dispatch idtype " << pytorch_dtype; \ + TORCH_CHECK(false, oss.str()); \ + return false; \ + } \ + }() + +#define _DISPATCH_CASE_F32(c_type, ...) \ + case at::ScalarType::Float: { \ + using c_type = float; \ + return __VA_ARGS__(); \ + } + +#ifdef FLASHINFER_ENABLE_F16 +#define _DISPATCH_CASE_F16(c_type, ...) \ + case at::ScalarType::Half: { \ + using c_type = nv_half; \ + return __VA_ARGS__(); \ + } +#else +#define _DISPATCH_CASE_F16(c_type, ...) +#endif + +#ifdef FLASHINFER_ENABLE_BF16 +#define _DISPATCH_CASE_BF16(c_type, ...) \ + case at::ScalarType::BFloat16: { \ + using c_type = nv_bfloat16; \ + return __VA_ARGS__(); \ + } +#else +#define _DISPATCH_CASE_BF16(c_type, ...) +#endif + +#ifdef FLASHINFER_ENABLE_FP8_E4M3 +#define _DISPATCH_CASE_FP8_E4M3(c_type, ...) \ + case at::ScalarType::Float8_e4m3fn: { \ + using c_type = __nv_fp8_e4m3; \ + return __VA_ARGS__(); \ + } +#else +#define _DISPATCH_CASE_FP8_E4M3(c_type, ...) +#endif + +#ifdef FLASHINFER_ENABLE_FP8_E5M2 +#define _DISPATCH_CASE_FP8_E5M2(c_type, ...) \ + case at::ScalarType::Float8_e5m2: { \ + using c_type = __nv_fp8_e5m2; \ + return __VA_ARGS__(); \ + } +#else +#define _DISPATCH_CASE_FP8_E5M2(c_type, ...) +#endif + +// Should not be used together with _DISPATCH_SF_CASE_FP8_E8M0 +#ifdef FLASHINFER_ENABLE_FP4_E2M1 +#if (__CUDACC_VER_MAJOR__ * 10000 + __CUDACC_VER_MINOR__ * 100 >= 120800) +#define _DISPATCH_CASE_FP4_E2M1(c_type, ...) \ + case at::ScalarType::Byte: { \ + using c_type = __nv_fp4_e2m1; \ + return __VA_ARGS__(); \ + } +#else +#define _DISPATCH_CASE_FP4_E2M1(c_type, ...) \ + case at::ScalarType::Byte: { \ + static_assert(false, "FP4 E2M1 support requires CUDA 12.8 or newer."); \ + break; \ + } +#endif +#else +#define _DISPATCH_CASE_FP4_E2M1(c_type, ...) +#endif + +// Should not be used together with _DISPATCH_CASE_FP4_E2M1 +#ifdef FLASHINFER_ENABLE_FP8_E8M0 +#if (__CUDACC_VER_MAJOR__ * 10000 + __CUDACC_VER_MINOR__ * 100 >= 120800) +#define _DISPATCH_SF_CASE_FP8_E8M0(c_type, ...) \ + case at::ScalarType::Byte: { \ + using c_type = __nv_fp8_e8m0; \ + return __VA_ARGS__(); \ + } +#else +#define _DISPATCH_SF_CASE_FP8_E8M0(c_type, ...) \ + case at::ScalarType::Byte: { \ + static_assert(false, "FP8 E8M0 support requires CUDA 12.8 or newer."); \ + break; \ + } +#endif +#else +#define _DISPATCH_SF_CASE_FP8_E8M0(c_type, ...) +#endif + +#define DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FP16(pytorch_dtype, c_type, ...) \ + [&]() -> bool { \ + switch (pytorch_dtype) { \ + _DISPATCH_CASE_F16(c_type, __VA_ARGS__) \ + _DISPATCH_CASE_BF16(c_type, __VA_ARGS__) \ + default: \ + std::ostringstream oss; \ + oss << __PRETTY_FUNCTION__ << " failed to dispatch data type " << pytorch_dtype; \ + TORCH_CHECK(false, oss.str()); \ + return false; \ + } \ + }() + +#define DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FP8(pytorch_dtype, c_type, ...) \ + [&]() -> bool { \ + switch (pytorch_dtype) { \ + _DISPATCH_CASE_FP8_E4M3(c_type, __VA_ARGS__) \ + _DISPATCH_CASE_FP8_E5M2(c_type, __VA_ARGS__) \ + default: \ + std::ostringstream oss; \ + oss << __PRETTY_FUNCTION__ << " failed to dispatch fp8 data type " << pytorch_dtype; \ + TORCH_CHECK(false, oss.str()); \ + return false; \ + } \ + }() + +#define DISPATCH_PYTORCH_DTYPE_TO_CTYPE_SF(pytorch_dtype, c_type, ...) \ + [&]() -> bool { \ + switch (pytorch_dtype) { \ + _DISPATCH_CASE_F32(c_type, __VA_ARGS__) \ + _DISPATCH_SF_CASE_FP8_E8M0(c_type, __VA_ARGS__) \ + default: \ + std::ostringstream oss; \ + oss << __PRETTY_FUNCTION__ << " failed to dispatch scaling factor data type " \ + << pytorch_dtype; \ + TORCH_CHECK(false, oss.str()); \ + return false; \ + } \ + }() + +#define DISPATCH_PYTORCH_DTYPE_TO_CTYPE(pytorch_dtype, c_type, ...) \ + [&]() -> bool { \ + switch (pytorch_dtype) { \ + _DISPATCH_CASE_F16(c_type, __VA_ARGS__) \ + _DISPATCH_CASE_BF16(c_type, __VA_ARGS__) \ + _DISPATCH_CASE_FP8_E4M3(c_type, __VA_ARGS__) \ + _DISPATCH_CASE_FP8_E5M2(c_type, __VA_ARGS__) \ + _DISPATCH_CASE_FP4_E2M1(c_type, __VA_ARGS__) \ + default: \ + std::ostringstream oss; \ + oss << __PRETTY_FUNCTION__ << " failed to dispatch data type " << pytorch_dtype; \ + TORCH_CHECK(false, oss.str()); \ + return false; \ + } \ + }() + +#define _DISPATCH_SWITCH(var_name, cond, ...) \ + [&]() -> bool { \ + switch (cond) { \ + __VA_ARGS__ \ + default: \ + std::ostringstream oss; \ + oss << __PRETTY_FUNCTION__ << " failed to dispatch " var_name " " << int(cond); \ + TORCH_CHECK(false, oss.str()); \ + return false; \ + } \ + }() + +#define _DISPATCH_SWITCH_U16x2(var1_name, var2_name, cond1, cond2, ...) \ + [&]() -> bool { \ + switch (pack_u16(cond1, cond2)) { \ + __VA_ARGS__ \ + default: \ + std::ostringstream oss; \ + oss << __PRETTY_FUNCTION__ << " failed to dispatch (" var1_name ", " var2_name "): (" \ + << int(cond1) << ", " << int(cond2) << ")"; \ + TORCH_CHECK(false, oss.str()); \ + return false; \ + } \ + }() + +#define _DISPATCH_CASE(case_expr, case_var, ...) \ + case case_expr: { \ + constexpr auto case_var = case_expr; \ + return __VA_ARGS__(); \ + } + +#define _DISPATCH_CASE_U16x2(case_expr1, case_expr2, case_var1, case_var2, ...) \ + case pack_u16(case_expr1, case_expr2): { \ + constexpr auto case_var1 = case_expr1; \ + constexpr auto case_var2 = case_expr2; \ + return __VA_ARGS__(); \ + } + +#define DISPATCH_BOOL(expr, const_expr, ...) \ + [&]() -> bool { \ + if (expr) { \ + constexpr bool const_expr = true; \ + return __VA_ARGS__(); \ + } else { \ + constexpr bool const_expr = false; \ + return __VA_ARGS__(); \ + } \ + }() + +inline void check_shape(const at::Tensor& a, const at::Tensor& b, const char* a_name, + const char* b_name) { + TORCH_CHECK(a.dim() == b.dim(), a_name, ".dim() != ", b_name, ".dim(). ", a.dim(), " vs ", + b.dim()); + for (int i = 0; i < a.dim(); ++i) { + TORCH_CHECK(a.size(i) == b.size(i), a_name, ".size(", i, ") != ", b_name, ".size(", i, ")"); + } +} + +inline constexpr uint32_t pack_u16(uint16_t a, uint16_t b) { + return (uint32_t(a) << 16) | uint32_t(b); +} + +#define CHECK_GQA_HEAD_DIVISIBLE(num_qo_heads, num_kv_heads) \ + TORCH_CHECK(num_qo_heads % num_kv_heads == 0, "num_qo_heads(", num_qo_heads, \ + ") must be divisible by num_kv_heads(", num_kv_heads, ")") + +#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor") + +#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous") +#define CHECK_LAST_DIM_CONTIGUOUS(x) \ + TORCH_CHECK(x.strides()[x.strides().size() - 1] == 1, #x "must be contiguous at last dimension") + +#define CHECK_INPUT(x) \ + CHECK_CUDA(x); \ + CHECK_CONTIGUOUS(x) +#define CHECK_INPUT_TYPE(x, st) \ + TORCH_CHECK(x.scalar_type() == st, "Inconsistency of Tensor type: " #x) +#define CHECK_INPUT_AND_TYPE(x, st) \ + CHECK_CUDA(x); \ + CHECK_CONTIGUOUS(x); \ + CHECK_INPUT_TYPE(x, st) +#define CHECK_LAST_DIM_CONTIGUOUS_INPUT(x) \ + CHECK_CUDA(x); \ + CHECK_LAST_DIM_CONTIGUOUS(x) + +#define CHECK_DIM(d, x) TORCH_CHECK(x.dim() == d, #x " must be a " #d "D tensor") + +#define CHECK_SHAPE(a, b) check_shape(a, b, #a, #b) + +#define CHECK_EQ(a, b) TORCH_CHECK((a) == (b), "CHECK_EQ(" #a ", " #b ") failed. ", a, " vs ", b) + +#define CHECK_GE(a, b) TORCH_CHECK((a) >= (b), "CHECK_GE(" #a ", " #b ") failed. ", a, " vs ", b) + +inline bool is_float8_tensor(const at::Tensor& tensor) { + return tensor.scalar_type() == at::ScalarType::Float8_e4m3fn || + tensor.scalar_type() == at::ScalarType::Float8_e5m2; +} diff --git a/tests/simple_test.py b/tests/simple_test.py new file mode 100644 index 0000000000000000000000000000000000000000..e061586e7c661490d0f7e43ded6d51169821e02e --- /dev/null +++ b/tests/simple_test.py @@ -0,0 +1,9 @@ +import torch +import flashinfer + +out = torch.empty(3, 3, dtype=torch.float16).cuda() +input_tensor = torch.randn(3, 3, dtype=torch.float16).cuda() +enable_pdl = True + +flashinfer.gelu_and_mul(out, input_tensor, enable_pdl) +print(out) \ No newline at end of file diff --git a/torch-ext/flashinfer/__init__.py b/torch-ext/flashinfer/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..077f829fb7146709620e08250a58523955e1486a --- /dev/null +++ b/torch-ext/flashinfer/__init__.py @@ -0,0 +1,49 @@ +# NOTE: Torch needs to be imported before the custom +# extensions. Otherwise libc10.so cannot be found. +import torch +from ._ops import ops + +def gelu_and_mul(out, input, enable_pdl): + return ops.gelu_and_mul(out, input, enable_pdl) + +def gelu_tanh_and_mul(out, input, enable_pdl): + return ops.gelu_tanh_and_mul(out, input, enable_pdl) + +def silu_and_mul(out, input, enable_pdl): + return ops.silu_and_mul(out, input, enable_pdl) + +# def batch_decode_with_paged_kv_cache_plan_swa_false( +# float_workspace_buffer, +# int_workspace_buffer, +# page_locked_int_workspace_buffer, +# indptr, +# batch_size, +# num_qo_heads, +# num_kv_heads, +# page_size, +# enable_cuda_graph, +# window_left, +# logits_soft_cap, +# head_dim_qk, +# head_dim_vo, +# empty_q_data, +# empty_kv_data +# ): +# return ops.BatchDecodeWithPagedKVCachePlan_swa_false( +# float_workspace_buffer, +# int_workspace_buffer, +# page_locked_int_workspace_buffer, +# indptr, +# batch_size, +# num_qo_heads, +# num_kv_heads, +# page_size, +# enable_cuda_graph, +# window_left, +# logits_soft_cap, +# head_dim_qk, +# head_dim_vo, +# empty_q_data, +# empty_kv_data +# ) + \ No newline at end of file diff --git a/torch-ext/torch_binding.cpp b/torch-ext/torch_binding.cpp new file mode 100644 index 0000000000000000000000000000000000000000..08de65bfee8c3b5472d39403f79e127445f169f1 --- /dev/null +++ b/torch-ext/torch_binding.cpp @@ -0,0 +1,107 @@ +#include + +#include "registration.h" +#include "torch_binding.h" + +TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { + + ops.def("gelu_and_mul(Tensor! out, Tensor! input, bool enable_pdl) -> ()"); + ops.impl("gelu_and_mul", torch::kCUDA, &gelu_and_mul); + + ops.def("gelu_tanh_and_mul(Tensor! out, Tensor! input, bool enable_pdl) -> ()"); + ops.impl("gelu_tanh_and_mul", torch::kCUDA, &gelu_tanh_and_mul); + + ops.def("silu_and_mul(Tensor! out, Tensor! input, bool enable_pdl) -> ()"); + ops.impl("silu_and_mul", torch::kCUDA, &silu_and_mul); + + // Register variant-specific operators directly + // ops.def( + // "BatchDecodeWithPagedKVCachePlan_swa_false(" + // "Tensor float_workspace_buffer, " + // "Tensor int_workspace_buffer, " + // "Tensor page_locked_int_workspace_buffer, " + // "Tensor indptr, " + // "int batch_size, " + // "int num_qo_heads, " + // "int num_kv_heads, " + // "int page_size, " + // "bool enable_cuda_graph, " + // "int window_left, " + // "float logits_soft_cap, " + // "int head_dim_qk, " + // "int head_dim_vo, " + // "Tensor empty_q_data, " + // "Tensor empty_kv_data) -> Tensor" + // ); + // ops.impl("BatchDecodeWithPagedKVCachePlan_swa_false", torch::kCUDA, &BatchDecodeWithPagedKVCachePlan_swa_false); + + // ops.def( + // "BatchDecodeWithPagedKVCacheRun_swa_false(" + // "Tensor float_workspace_buffer, " + // "Tensor int_workspace_buffer, " + // "Tensor plan_info_vec, " + // "Tensor q, " + // "Tensor paged_k_cache, " + // "Tensor paged_v_cache, " + // "Tensor paged_kv_indptr, " + // "Tensor paged_kv_indices, " + // "Tensor paged_kv_last_page_len, " + // "Tensor o, " + // "Tensor? maybe_lse, " + // "int kv_layout_code, " + // "int window_left, " + // "bool enable_pdl, " + // "Tensor? maybe_alibi_slopes, " + // "float logits_soft_cap, " + // "float sm_scale, " + // "float rope_rcp_scale, " + // "float rope_rcp_theta) -> ()" + // ); + // ops.impl("BatchDecodeWithPagedKVCacheRun_swa_false", torch::kCUDA, &BatchDecodeWithPagedKVCacheRun_swa_false); + + // ops.def( + // "BatchDecodeWithPagedKVCachePlan_swa_true(" + // "Tensor float_workspace_buffer, " + // "Tensor int_workspace_buffer, " + // "Tensor page_locked_int_workspace_buffer, " + // "Tensor indptr, " + // "int batch_size, " + // "int num_qo_heads, " + // "int num_kv_heads, " + // "int page_size, " + // "bool enable_cuda_graph, " + // "int window_left, " + // "float logits_soft_cap, " + // "int head_dim_qk, " + // "int head_dim_vo, " + // "Tensor empty_q_data, " + // "Tensor empty_kv_data) -> Tensor" + // ); + // ops.impl("BatchDecodeWithPagedKVCachePlan_swa_true", torch::kCUDA, &BatchDecodeWithPagedKVCachePlan_swa_true); + + // ops.def( + // "BatchDecodeWithPagedKVCacheRun_swa_true(" + // "Tensor float_workspace_buffer, " + // "Tensor int_workspace_buffer, " + // "Tensor plan_info_vec, " + // "Tensor q, " + // "Tensor paged_k_cache, " + // "Tensor paged_v_cache, " + // "Tensor paged_kv_indptr, " + // "Tensor paged_kv_indices, " + // "Tensor paged_kv_last_page_len, " + // "Tensor o, " + // "Tensor? maybe_lse, " + // "int kv_layout_code, " + // "int window_left, " + // "bool enable_pdl, " + // "Tensor? maybe_alibi_slopes, " + // "float logits_soft_cap, " + // "float sm_scale, " + // "float rope_rcp_scale, " + // "float rope_rcp_theta) -> ()" + // ); + // ops.impl("BatchDecodeWithPagedKVCacheRun_swa_true", torch::kCUDA, &BatchDecodeWithPagedKVCacheRun_swa_true); +} + +REGISTER_EXTENSION(TORCH_EXTENSION_NAME) diff --git a/torch-ext/torch_binding.h b/torch-ext/torch_binding.h new file mode 100644 index 0000000000000000000000000000000000000000..a0f9405632935ff479793f470ff5be3442b85bd6 --- /dev/null +++ b/torch-ext/torch_binding.h @@ -0,0 +1,101 @@ +#pragma once + +#include + +// at::Tensor BatchDecodeWithPagedKVCachePlan_swa_false( +// at::Tensor float_workspace_buffer, +// at::Tensor int_workspace_buffer, +// at::Tensor page_locked_int_workspace_buffer, +// at::Tensor indptr, +// int64_t batch_size, +// int64_t num_qo_heads, +// int64_t num_kv_heads, +// int64_t page_size, +// bool enable_cuda_graph, +// int64_t window_left, +// double logits_soft_cap, +// int64_t head_dim_qk, +// int64_t head_dim_vo, +// at::Tensor empty_q_data, +// at::Tensor empty_kv_data +// ); + +// void BatchDecodeWithPagedKVCacheRun_swa_false( +// at::Tensor float_workspace_buffer, +// at::Tensor int_workspace_buffer, +// at::Tensor plan_info_vec, +// at::Tensor q, +// at::Tensor paged_k_cache, +// at::Tensor paged_v_cache, +// at::Tensor paged_kv_indptr, +// at::Tensor paged_kv_indices, +// at::Tensor paged_kv_last_page_len, +// at::Tensor o, +// std::optional maybe_lse, +// int64_t kv_layout_code, +// int64_t window_left, +// bool enable_pdl, +// std::optional maybe_alibi_slopes, +// double logits_soft_cap, +// double sm_scale, +// double rope_rcp_scale, +// double rope_rcp_theta +// ); + +// at::Tensor BatchDecodeWithPagedKVCachePlan_swa_true( +// at::Tensor float_workspace_buffer, +// at::Tensor int_workspace_buffer, +// at::Tensor page_locked_int_workspace_buffer, +// at::Tensor indptr, +// int64_t batch_size, +// int64_t num_qo_heads, +// int64_t num_kv_heads, +// int64_t page_size, +// bool enable_cuda_graph, +// int64_t window_left, +// double logits_soft_cap, +// int64_t head_dim_qk, +// int64_t head_dim_vo, +// at::Tensor empty_q_data, +// at::Tensor empty_kv_data +// ); + +// void BatchDecodeWithPagedKVCacheRun_swa_true( +// at::Tensor float_workspace_buffer, +// at::Tensor int_workspace_buffer, +// at::Tensor plan_info_vec, +// at::Tensor q, +// at::Tensor paged_k_cache, +// at::Tensor paged_v_cache, +// at::Tensor paged_kv_indptr, +// at::Tensor paged_kv_indices, +// at::Tensor paged_kv_last_page_len, +// at::Tensor o, +// std::optional maybe_lse, +// int64_t kv_layout_code, +// int64_t window_left, +// bool enable_pdl, +// std::optional maybe_alibi_slopes, +// double logits_soft_cap, +// double sm_scale, +// double rope_rcp_scale, +// double rope_rcp_theta +// ); + +void gelu_and_mul( + at::Tensor& out, + at::Tensor& input, + bool enable_pdl +); + +void gelu_tanh_and_mul( + at::Tensor& out, + at::Tensor& input, + bool enable_pdl +); + +void silu_and_mul( + at::Tensor& out, + at::Tensor& input, + bool enable_pdl +); \ No newline at end of file